Added bz2 file writing support to close #73

functional/io.py

@@ -214,6 +214,9 @@ def universal_write_open(path, mode, buffering=-1, encoding=None, errors=None, n
     elif compression == 'lzma' or compression == 'xz':
         return lzma.open(path, mode=mode, format=format, check=check, preset=preset,
                          filters=filters, encoding=encoding, errors=errors, newline=newline)
+    elif compression == 'bz2':
+        return bz2.open(path, mode=mode, compresslevel=compresslevel, encoding=encoding,
+                        errors=errors, newline=newline)
     else:
         raise ValueError(
             'compression must be None, gz, gzip, lzma, or xz and was {0}'.format(compression))

functional/test/test_streams.py

@@ -10,7 +10,7 @@
 
 from functional import seq, pseq
 from functional.streams import Stream, ParallelStream
-from functional.io import lzma
+from functional.io import lzma, bz2
 
 
 class TestStreams(unittest.TestCase):
@@ -237,6 +237,10 @@ def test_to_file_compressed(self):
         with lzma.open(tmp_path, 'rt') as output:
             self.assertEqual('[1, 2, 3, 4]', output.readlines()[0])
 
+        sequence.to_file(tmp_path, compression='bz2')
+        with bz2.open(tmp_path, 'rt') as output:
+            self.assertEqual('[1, 2, 3, 4]', output.readlines()[0])
+
     def test_to_jsonl(self):
         tmp_path = 'functional/test/data/tmp/output.txt'
         elements = [{'a': 1, 'b': 2}, {'c': 3}, {'d': 4}]
@@ -259,6 +263,10 @@ def test_to_jsonl_compressed(self):
         result = self.seq.jsonl(tmp_path).to_list()
         self.assertEqual(elements, result)
 
+        sequence.to_jsonl(tmp_path, compression='bz2')
+        result = self.seq.jsonl(tmp_path).to_list()
+        self.assertEqual(elements, result)
+
     def test_to_json(self):
         tmp_path = 'functional/test/data/tmp/output.txt'
         elements = [[u'a', 1], [u'b', 2], [u'c', 3]]
@@ -276,13 +284,13 @@ def test_to_json(self):
     def test_to_json_compressed(self):
         tmp_path = 'functional/test/data/tmp/output.txt'
         elements = [[u'a', 1], [u'b', 2], [u'c', 3]]
+        dict_expect = {u'a': 1, u'b': 2, u'c': 3}
         sequence = self.seq(elements)
 
         sequence.to_json(tmp_path, compression='gzip')
         result = self.seq.json(tmp_path).to_list()
         self.assertEqual(elements, result)
 
-        dict_expect = {u'a': 1, u'b': 2, u'c': 3}
         sequence.to_json(tmp_path, root_array=False, compression='gzip')
         result = self.seq.json(tmp_path).to_dict()
         self.assertEqual(dict_expect, result)
@@ -291,11 +299,18 @@ def test_to_json_compressed(self):
         result = self.seq.json(tmp_path).to_list()
         self.assertEqual(elements, result)
 
-        dict_expect = {u'a': 1, u'b': 2, u'c': 3}
         sequence.to_json(tmp_path, root_array=False, compression='lzma')
         result = self.seq.json(tmp_path).to_dict()
         self.assertEqual(dict_expect, result)
 
+        sequence.to_json(tmp_path, compression='bz2')
+        result = self.seq.json(tmp_path).to_list()
+        self.assertEqual(elements, result)
+
+        sequence.to_json(tmp_path, root_array=False, compression='bz2')
+        result = self.seq.json(tmp_path).to_dict()
+        self.assertEqual(dict_expect, result)
+
     def test_to_csv(self):
         tmp_path = 'functional/test/data/tmp/output.txt'
         elements = [[1, 2, 3], [4, 5, 6], ['a', 'b', 'c']]
@@ -319,6 +334,10 @@ def test_to_csv_compressed(self):
         result = self.seq.csv(tmp_path).to_list()
         self.assertEqual(expect, result)
 
+        sequence.to_csv(tmp_path, compression='bz2')
+        result = self.seq.csv(tmp_path).to_list()
+        self.assertEqual(expect, result)
+
     def test_to_sqlite3_failure(self):
         insert_sql = 'INSERT INTO user (id, name) VALUES (?, ?)'
         elements = [(1, 'Tom'), (2, 'Jack'), (3, 'Jane'), (4, 'Stephan')]