From d9c5f0d066833bec3eee680a6e50df93fa4d750e Mon Sep 17 00:00:00 2001
From: ChillerDragon <ChillerDragon@gmail.com>
Date: Sun, 18 Jun 2023 12:51:07 +0200
Subject: [PATCH] feat: add optional libtw2 huffman backend

Prefer libtw2 over TeeAI huffman but do not install it by default.
Thanks to @heinrich5991 for porting his rust code to python.
---
 README.md                 | 23 +++++++++++++++++++++++
 requirements/dev.txt      |  1 +
 requirements/optional.txt |  1 +
 twnet_parser/huffman.py   | 24 ++++++++++++++++++++++++
 twnet_parser/packet.py    |  6 +++---
 5 files changed, 52 insertions(+), 3 deletions(-)
 create mode 100644 requirements/optional.txt
 create mode 100644 twnet_parser/huffman.py

diff --git a/README.md b/README.md
index 50b664f..82b43e0 100644
--- a/README.md
+++ b/README.md
@@ -125,6 +125,29 @@ packet.messages.append(msg)
 packet.pack() # => b'\x02~\x06H\x1f\x93\xd7\x00\x00\x80\x01option\x00test\x00\x00\x00'
 ```
 
+## Zero dependencies by default
+
+Running ``pip install twnet_parser`` will not install any additional packages.
+
+But there is an optional dependency for huffman compression.
+By default twnet_parser is using the huffman compression code from the [TeeAI](https://github.com/edg-l/TeeAI)
+project which is written in pure python.
+If you have [libtw2-huffman](https://pypi.org/project/libtw2-huffman/) installed it will use that one instead.
+Because it is faster since it is written in rust and has better error handling.
+But since it is so much overhead it is not installed by default to keep twnet_parser light weight.
+
+
+You can install it by running ``pip install libtw2-huffman``
+or by running ``pip install -r requirements/optional.txt``
+
+
+You can also check which huffman backend is currently active with these lines of code
+
+```python
+import twnet_parser.huffman
+print(twnet_parser.huffman.backend_name()) # => rust-libtw2 or python-TeeAI
+```
+
 ## development setup
 
 ```bash
diff --git a/requirements/dev.txt b/requirements/dev.txt
index e66e471..fed4c77 100644
--- a/requirements/dev.txt
+++ b/requirements/dev.txt
@@ -1,4 +1,5 @@
 -r prod.txt
+-r optional.txt
 astroid==2.15.0
 attrs==22.2.0
 bleach==6.0.0
diff --git a/requirements/optional.txt b/requirements/optional.txt
new file mode 100644
index 0000000..23cfb96
--- /dev/null
+++ b/requirements/optional.txt
@@ -0,0 +1 @@
+libtw2-huffman==0.2.0
diff --git a/twnet_parser/huffman.py b/twnet_parser/huffman.py
new file mode 100644
index 0000000..4374dd8
--- /dev/null
+++ b/twnet_parser/huffman.py
@@ -0,0 +1,24 @@
+from twnet_parser.external.huffman import huffman
+
+HAS_LIBTW2 = False
+
+try:
+    import libtw2_huffman # type: ignore
+    HAS_LIBTW2 = True
+except ImportError:
+    HAS_LIBTW2 = False
+
+def backend_name() -> str:
+    if HAS_LIBTW2:
+        return 'rust-libtw2'
+    return 'python-TeeAI'
+
+def compress(data: bytes) -> bytes:
+    if HAS_LIBTW2:
+        return libtw2_huffman.compress(data)
+    return huffman.decompress(bytearray(data))
+
+def decompress(data: bytes) -> bytes:
+    if HAS_LIBTW2:
+        return libtw2_huffman.decompress(data)
+    return huffman.decompress(bytearray(data))
diff --git a/twnet_parser/packet.py b/twnet_parser/packet.py
index e02fd7c..7fb5ebd 100644
--- a/twnet_parser/packet.py
+++ b/twnet_parser/packet.py
@@ -15,7 +15,7 @@ from twnet_parser.msg_matcher.control6 import match_control6
 from twnet_parser.msg_matcher.connless6 import match_connless6
 from twnet_parser.constants import NET_MAX_SEQUENCE, NET_PACKETVERSION
 
-from twnet_parser.external.huffman import huffman
+import twnet_parser.huffman
 
 # TODO: what is a nice pythonic way of storing those?
 #       also does some version:: namespace thing make sense?
@@ -419,7 +419,7 @@ class PacketParser():
             return pck
         if pck.header.flags.compression:
             payload = bytearray(pck.payload_raw)
-            pck.payload_decompressed = huffman.decompress(payload)
+            pck.payload_decompressed = twnet_parser.huffman.decompress(payload)
         pck.messages = cast(
                 list[Union[CtrlMessage, NetMessage, ConnlessMessage]],
                 self.get_messages(pck.payload_decompressed))
@@ -448,7 +448,7 @@ class PacketParser():
             return pck
         if pck.header.flags.compression:
             payload = bytearray(pck.payload_raw)
-            pck.payload_decompressed = huffman.decompress(payload)
+            pck.payload_decompressed = twnet_parser.huffman.decompress(payload)
         pck.messages = cast(
                 list[Union[CtrlMessage, NetMessage, ConnlessMessage]],
                 self.get_messages(pck.payload_decompressed))