Add support for zstd compression algorithm

.github/workflows/test.yml

@@ -0,0 +1,26 @@
+name: test
+on:
+  pull_request:
+    branches-ignore:
+      - "release-**"
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ['3.7', '3.8', '3.9']
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+          check-latest: true
+      - name: Install dependencies
+        run: |
+          sudo apt-get install -y zstd
+          pip install -e .[test]
+      - name: Test with pytest
+        run: |
+          py.test -v -s tests

pylogrotate/main.py

@@ -36,6 +36,7 @@
     'group': 'root',
     # 'rotate': 7,
     'compress': True,
+    'compressalgorithm': 'gzip',
     'copy': [],
     'copytohdfs': [],
     'hdfs': {},
@@ -105,6 +106,12 @@ def gzip(path):
     run('gzip -kf {}'.format(path))
 
 
+def zstd(path):
+    if not path:
+        return
+    run('zstd -kf {}'.format(path))
+
+
 class Rotator(object):
 
     def __init__(self, config):
@@ -117,6 +124,7 @@ def __init__(self, config):
         # FIXME: Handle rotated files keeping correctly
         # self.keep_files = int(config['rotate'])
         self.compress = config['compress']
+        self.compress_algorithm = config['compressalgorithm']
 
         self.copy = config['copy']
         self.copytohdfs = config['copytohdfs']
@@ -172,8 +180,14 @@ def rename_file(self, path):
         return dest_path
 
     def compress_file(self, dest_path):
-        gzip(dest_path)
-        return '{}.gz'.format(dest_path)
+        if self.compress_algorithm == 'gzip':
+            gzip(dest_path)
+            return '{}.gz'.format(dest_path)
+        elif self.compress_algorithm == 'zstd':
+            zstd(dest_path)
+            return '{}.zst'.format(dest_path)
+        else:
+            raise ValueError('Unsupported compression algorithm: {}'.format(self.compress_algorithm))
 
     def _copy_file(self, path, from_, to):
         if not to:

setup.py

@@ -15,7 +15,7 @@
     url='https://github.com/xiachufang/pylogrotate',
     packages=find_packages(exclude=['contrib', 'docs', 'tests*']),
     install_requires=install_requires,
-    extras_require=dict(test=['pytest', 'freezegun']),
+    extras_require=dict(test=['pytest', 'freezegun', 'zstandard']),
     license='MIT',
     classifiers=[
         'Development Status :: 3 - Alpha',

tests/test_pylogrotate.py

@@ -2,6 +2,7 @@
 
 import grp
 import gzip
+import zstandard
 import io
 import os
 import pwd
@@ -83,6 +84,47 @@ def test_rotate(tmpdir):
         assert rf.read() == content
 
 
+def test_compress_gzip(tmpdir):
+    content = b'This log file should be compressed.'
+
+    f = tmpdir.mkdir('nginx').join('access.rotate-this.log')
+    f.write(content)
+
+    assert f.exists()
+    assert len(tmpdir.join('nginx').listdir()) == 1
+    rotator = Rotator(get_config(paths=[str(f)], compress=True, compressalgorithm='gzip'))
+    rotator.rotate()
+
+    # `f` compressed and moved into `*-rotate/` directory
+    assert not f.exists()
+    assert len(tmpdir.join('nginx').listdir()) == 1
+
+    rf = _traverse(tmpdir.join('nginx'))
+    with gzip.open(str(rf)) as rf:
+        assert rf.read() == content
+
+
+def test_compress_zstd(tmpdir):
+    content = b'This log file should be compressed.'
+
+    f = tmpdir.mkdir('nginx').join('access.rotate-this.log')
+    f.write(content)
+
+    assert f.exists()
+    assert len(tmpdir.join('nginx').listdir()) == 1
+    rotator = Rotator(get_config(paths=[str(f)], compress=True, compressalgorithm='zstd'))
+    rotator.rotate()
+
+    # `f` compressed and moved into `*-rotate/` directory
+    assert not f.exists()
+    assert len(tmpdir.join('nginx').listdir()) == 1
+
+    rf = _traverse(tmpdir.join('nginx'))
+    with open(str(rf), 'rb') as rf:
+        dctx = zstandard.ZstdDecompressor()
+        assert dctx.decompress(rf.read()) == content
+
+
 def test_skip_empty_files(tmpdir):
     f = tmpdir.mkdir('nginx').join('access.skip-empty.log')
     f.write('')