From 1b6a3580cc25d804089bc9b31faf7ddf6219b822 Mon Sep 17 00:00:00 2001
From: William Toohey <will@mon.im>
Date: Sun, 28 Apr 2024 10:54:05 +1000
Subject: [PATCH] Fix BrokenPipeError when piping kbinxml output

---
 kbinxml/kbinxml.py | 18 +++++++++++-------
 setup.py           |  3 +--
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/kbinxml/kbinxml.py b/kbinxml/kbinxml.py
index 37bde67..7f1a520 100644
--- a/kbinxml/kbinxml.py
+++ b/kbinxml/kbinxml.py
@@ -1,14 +1,14 @@
 import argparse
-from struct import calcsize
-import sys
 import operator
+import sys
 from io import BytesIO
+from struct import calcsize
 
 import lxml.etree as etree
 
 from .bytebuffer import ByteBuffer
-from .sixbit import pack_sixbit, unpack_sixbit
 from .format_ids import xml_formats, xml_types
+from .sixbit import pack_sixbit, unpack_sixbit
 
 DEBUG_OFFSETS = False
 DEBUG = False
@@ -479,10 +479,14 @@ def main():
 
     xml = KBinXML(input, convert_illegal_things=args.convert_illegal)
     stdout = getattr(sys.stdout, "buffer", sys.stdout)
-    if KBinXML.is_binary_xml(input):
-        stdout.write(xml.to_text().encode("utf-8"))
-    else:
-        stdout.write(xml.to_binary())
+    try:
+        if KBinXML.is_binary_xml(input):
+            stdout.write(xml.to_text().encode("utf-8"))
+        else:
+            stdout.write(xml.to_binary())
+    except BrokenPipeError:
+        # allows kbinxml to be piped to `head` or similar
+        sys.exit(141)
 
 
 if __name__ == "__main__":
diff --git a/setup.py b/setup.py
index a943199..f6590f9 100644
--- a/setup.py
+++ b/setup.py
@@ -1,13 +1,12 @@
 from setuptools import setup
 
-
 requires = [
     "lxml",
 ]
 
 python_requires = ">=3.10"
 
-version = "2.0"
+version = "2.1"
 setup(
     name="kbinxml",
     description="Decoder/encoder for Konami's binary XML format",