Commits

AJ R  committed 346cc54

Added unittest test_SavWriter_zsav.py. Tests pass in Python 2 and 3

  • Participants
  • Parent commits 4547290

Comments (0)

Files changed (2)

File savReaderWriter/savWriter.py

 from py3k import *
 from header import *
 if cWriterowOK:
-    cWriterow = cWriterow.cWriterow
+    cWriterow = cWriterow.cWriterow  # TODO: make this Python3-ready (or make a Python 3 version)
 
 class SavWriter(Header):
 
         """ This function writes one record, which is a Python list,
         compare this Python version with the Cython version cWriterow."""
         float_ = float
-        bytify = self.bytify
         encoding = self.fileEncoding
         for i, value in enumerate(record):
             varName = self.varNames[i]
                 except (ValueError, TypeError):
                     value = self.sysmis_
             else:
-                # Get rid of trailing null bytes --> 7 x faster than 'ljust'
-                #value = self.pad_8_lookup[varType] % value
-                template = self.pad_8_lookup[varType]
-                value = template % value.decode(encoding)  # TODO make this more efficient
-                value = value.encode(encoding)
+                if isPy3k:
+                    template = self.pad_8_lookup[varType]
+                    value = template % value.decode(encoding)  # TODO make this more efficient
+                    value = value.encode(encoding)
+                else:
+                    # Get rid of trailing null bytes --> 7 x faster than 'ljust'
+                    value = self.pad_8_lookup[varType] % value
+
                 if self.ioUtf8_ and isinstance(value, unicode):
                     value = value.encode("utf-8")          # TODO correct this
             record[i] = value

File savReaderWriter/unit_tests/test_SavWriter_zsav.py

+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+##############################################################################
+## Write a file using zlib compression (.zsav)
+##############################################################################
+
+import unittest
+import os
+import tempfile
+from savReaderWriter import *
+
+# .zsav = zlib compressed (requires IO libraries >= v21)
+# .zsav files CANNOT be read using SPSS v20 or lower
+class test_SavWriter_zsav(unittest.TestCase):
+    """Write a file using zlib compression (.zsav)"""
+
+    def setUp(self):
+        self.savFileName =  os.path.join(tempfile.gettempdir(), "test.zsav")
+        varNames = ['var1', 'v2', 'v3', 'bdate']
+        varTypes = {'var1': 6, 'v2': 0, 'v3': 0, 'bdate': 10}
+        self.args = (self.savFileName, varNames, varTypes)
+
+    def test_SavWriter_typical(self):
+        records_in = [[b'Test1', 1, 1, b'2010-08-11'],
+                      [b'Test2', 2, 1, b'1910-01-12']]
+        with SavWriter(*self.args) as writer:
+            for record in records_in:
+                writer.writerow(record)
+
+        reader = SavReader(self.savFileName)
+        with reader:  
+            records_out = [line for line in iter(reader)]
+            compression_got = reader.fileCompression
+        self.assertEqual(records_in, records_out)
+        self.assertEqual(b"zlib", compression_got)
+
+    def tearDown(self):
+        os.remove(self.savFileName)
+
+if __name__ == "__main__":
+    unittest.main()