dffml: sources: file: Added bz2 support

* dffml: sources: file: Added bz2 support
* tests: source: test_file: Added tests for bz2

Signed-off-by: John Andersen <[email protected]>

CHANGELOG.md

@@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - New Model creation tutorial
 - Added update functionality to the CSV source
 - Added support for Gzip file source
+- Added support for bz2 file source
 ### Changed
 - Restructured documentation to docs folder and moved from rST to markdown
 - Git feature cloc logs if no binaries are in path

dffml/source/file.py

@@ -3,7 +3,8 @@
 import os
 import abc
 import asyncio
-import gzip
+import gzip 
+import bz2
 from .source import Source
 from .log import LOGGER
 
@@ -42,6 +43,8 @@ async def _open(self):
             return
         if self.filename[::-1].startswith(('.gz')[::-1]):
             opener = gzip.open(self.filename, 'rt')
+        elif self.filename[::-1].startswith(('.bz2')[::-1]):
+            opener = bz2.open(self.filename, 'rt')
         else:
             opener = open(self.filename, 'r')
         with opener as fd:
@@ -54,6 +57,8 @@ async def _close(self):
         if not self.readonly:
             if self.filename[::-1].startswith(('.gz')[::-1]):
                 close = gzip.open(self.filename, 'wt')
+            elif self.filename[::-1].startswith(('.bz2')[::-1]):
+                close = bz2.open(self.filename, 'wt')
             else:
                 close = open(self.filename, 'w')
             with close as fd:

tests/source/test_file.py

@@ -71,6 +71,14 @@ async def test_open_gz(self):
             await source.open()
             m_open.assert_called_once_with('testfile.gz', 'rt')
 
+    async def test_open_bz2(self):
+        source = FakeFileSource('testfile.bz2')
+        m_open = mock_open()
+        with patch('os.path.exists', return_value=True), \
+                patch('bz2.open', m_open):
+            await source.open()
+            m_open.assert_called_once_with('testfile.bz2', 'rt')
+
     async def test_open_no_file(self):
         source = FakeFileSource('testfile')
         with patch('os.path.isfile', return_value=False):
@@ -91,6 +99,13 @@ async def test_close_gz(self):
             await source.close()
             m_open.assert_called_once_with('testfile.gz', 'wt')
 
+    async def test_close_bz2(self):
+        source = FakeFileSource('testfile.bz2')
+        m_open = mock_open()
+        with patch('bz2.open', m_open):
+            await source.close()
+            m_open.assert_called_once_with('testfile.bz2', 'wt')
+
     async def test_close_readonly(self):
         source = FakeFileSource('testfile:ro')
         m_open = mock_open()