From bbc470fc35013effdbe241a23da30b4a30d49afb Mon Sep 17 00:00:00 2001
From: meganomic <10358690+meganomic@users.noreply.github.com>
Date: Thu, 17 Dec 2020 11:20:03 +0100
Subject: [PATCH] Add support for lz4 files (#9)

---
 .github/workflows/ci.yml |  2 +-
 setup.py                 |  1 +
 xtarfile/lz4.py          | 34 ++++++++++++++++++++++++++++++++++
 xtarfile/xtarfile.py     |  2 ++
 4 files changed, 38 insertions(+), 1 deletion(-)
 create mode 100644 xtarfile/lz4.py

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index e27b76c..b1c1a25 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -32,5 +32,5 @@ jobs:
       - uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python }}
-      - run: pip install -e .[zstd]
+      - run: pip install -e .[lz4,zstd]
       - run: python setup.py test
diff --git a/setup.py b/setup.py
index ae134b7..ba4ee40 100644
--- a/setup.py
+++ b/setup.py
@@ -25,6 +25,7 @@
                 'compression formats.',
     long_description=long_description,
     extras_require={
+        'lz4': ['lz4 >= 2.2.1'],
         'zstd': ['zstandard >= 0.10.2']
     },
     python_requires='>=3.4',
diff --git a/xtarfile/lz4.py b/xtarfile/lz4.py
new file mode 100644
index 0000000..7dc4f59
--- /dev/null
+++ b/xtarfile/lz4.py
@@ -0,0 +1,34 @@
+from contextlib import contextmanager
+from tarfile import open as tarfile_open
+
+try:
+    import lz4.frame as lz4
+except ImportError:
+    lz4 = None
+
+
+class Lz4Tarfile:
+    def __init__(self, **kwargs):
+        self.lz4_kwargs = kwargs
+
+    @contextmanager
+    def read(self, path: str, mode: str):
+        with lz4.LZ4FrameFile(path) as lz4d:
+            archive = tarfile_open(mode=mode, fileobj=lz4d, **self.lz4_kwargs)
+            try:
+                yield archive
+            finally:
+                archive.close()
+
+    @contextmanager
+    def write(self, path: str, mode: str):
+        with lz4.LZ4FrameFile(path, mode=mode[0]) as lz4c:
+            archive = tarfile_open(mode=mode, fileobj=lz4c, **self.lz4_kwargs)
+            try:
+                yield archive
+            finally:
+                archive.close()
+
+
+if lz4 is None:
+    Lz4Tarfile = None  # noqa F811
diff --git a/xtarfile/xtarfile.py b/xtarfile/xtarfile.py
index 54aa0b9..4baa626 100644
--- a/xtarfile/xtarfile.py
+++ b/xtarfile/xtarfile.py
@@ -2,11 +2,13 @@
 from tarfile import open as tarfile_open
 
 from xtarfile.zstd import ZstandardTarfile
+from xtarfile.lz4 import Lz4Tarfile
 
 
 _HANDLERS = {
     'zstd': ZstandardTarfile,
     'zst': ZstandardTarfile,
+    'lz4': Lz4Tarfile,
 }
 
 _NATIVE_FORMATS = ('gz', 'bz2', 'xz', 'tar')