fix Surface.read/write_triangular(): parse error / invalid date str when locale != en

freesurfer_surface/__init__.py

@@ -20,7 +20,9 @@ https://surfer.nmr.mgh.harvard.edu/
 """
 
 import collections
+import contextlib
 import datetime
+import locale
 import re
 import struct
 import typing
@@ -30,6 +32,23 @@ try:
 except ImportError:  # pragma: no cover
     __version__ = None
 
+
+class UnsupportedLocaleSettingError(locale.Error):
+    pass
+
+
+@contextlib.contextmanager
+def setlocale(temporary_locale):
+    primary_locale = locale.setlocale(locale.LC_ALL)
+    try:
+        yield locale.setlocale(locale.LC_ALL, temporary_locale)
+    except locale.Error as exc:
+        if str(exc) == 'unsupported locale setting':
+            raise UnsupportedLocaleSettingError(temporary_locale)
+        raise exc
+    finally:
+        locale.setlocale(locale.LC_ALL, primary_locale)
+
 Vertex = collections.namedtuple('Vertex', ['right', 'anterior', 'superior'])
 
 
@@ -41,7 +60,6 @@ class Surface:
     _TAG_OLD_SURF_GEOM = b'\x00\x00\x00\x14'
     _TAG_OLD_USEREALRAS = b'\x00\x00\x00\x02'
 
-    # TODO set locale
     _DATETIME_FORMAT = '%a %b %d %H:%M:%S %Y'
 
     def __init__(self):
@@ -49,7 +67,6 @@ class Surface:
         self.creation_datetime = None
         self.vertices = []
         self.triangles_vertex_indices = []
-        # self._triangles = []
         self.using_old_real_ras = False
         self.volume_geometry_info = None
         self.command_lines = []
@@ -71,8 +88,9 @@ class Surface:
         assert stream.read(3) == self._MAGIC_NUMBER
         self.creator, creation_dt_str = re.match(rb'^created by (\w+) on (.* \d{4})\n',
                                                  stream.readline()).groups()
-        self.creation_datetime = datetime.datetime.strptime(creation_dt_str.decode(),
-                                                            self._DATETIME_FORMAT)
+        with setlocale('C'):
+            self.creation_datetime = datetime.datetime.strptime(creation_dt_str.decode(),
+                                                                self._DATETIME_FORMAT)
         assert stream.read(1) == b'\n'
         # fwriteInt
         # https://github.com/freesurfer/freesurfer/blob/release_6_0_0/utils/fio.c#L290
@@ -104,12 +122,15 @@ class Surface:
 
     def _triangular_creation_datetime_strftime(self) -> bytes:
         fmt = self._DATETIME_FORMAT.replace('%d', '{:>2}'.format(self.creation_datetime.day))
-        return self.creation_datetime.strftime(fmt).encode()
+        with setlocale('C'):
+            return self.creation_datetime.strftime(fmt).encode()
 
     def write_triangular(self, surface_file_path: str,
                          creation_datetime: typing.Optional[datetime.datetime] = None):
         if creation_datetime is None:
-            creation_datetime = datetime.datetime.now()
+            self.creation_datetime = datetime.datetime.now()
+        else:
+            self.creation_datetime = creation_datetime
         with open(surface_file_path, 'wb') as surface_file:
             surface_file.write(
                 self._MAGIC_NUMBER

tests/test_setlocale.py

@@ -0,0 +1,21 @@
+import locale
+
+import pytest
+
+from freesurfer_surface import setlocale, UnsupportedLocaleSettingError
+
+
+def test_set():
+    system_locale = locale.setlocale(locale.LC_ALL)
+    assert system_locale != 'C'
+    with setlocale('C'):
+        assert locale.setlocale(locale.LC_ALL) == 'C'
+    assert locale.setlocale(locale.LC_ALL) == system_locale
+
+
+def test_invalid():
+    system_locale = locale.setlocale(locale.LC_ALL)
+    with pytest.raises(UnsupportedLocaleSettingError):
+        with setlocale('abcdef21'):
+            pass
+    assert locale.setlocale(locale.LC_ALL) == system_locale

tests/test_surface.py

@@ -3,7 +3,7 @@ import os
 
 import pytest
 
-from freesurfer_surface import Surface, Vertex
+from freesurfer_surface import setlocale, Surface, Vertex
 
 
 SUBJECTS_DIR = os.path.join(os.path.dirname(__file__), 'subjects')
@@ -54,6 +54,12 @@ def test_read_triangular():
         b'  ']
 
 
+def test_read_triangular_locale():
+    with setlocale('de_AT.utf8'):
+        surface = Surface.read_triangular(SURFACE_FILE_PATH)
+    assert surface.creation_datetime == datetime.datetime(2019, 5, 9, 22, 37, 41)
+
+
 @pytest.mark.parametrize(('creation_datetime', 'expected_str'), [
     (datetime.datetime(2019, 5, 9, 22, 37, 41), b'Thu May  9 22:37:41 2019'),
     (datetime.datetime(2019, 4, 24, 23, 29, 22), b'Wed Apr 24 23:29:22 2019'),
@@ -96,6 +102,20 @@ def test_write_read_triangular_same(tmpdir):
     assert vars(expected_surface) == vars(resulted_surface)
 
 
+def test_write_triangular_same_locale(tmpdir):
+    surface = Surface()
+    surface.creator = b'pytest'
+    surface.volume_geometry_info = tuple(b'?' for _ in range(8))
+    creation_datetime = datetime.datetime(2018, 12, 31, 21, 42)
+    output_path = tmpdir.join('surface')
+    with setlocale('de_AT.utf8'):
+        surface.write_triangular(output_path, creation_datetime=creation_datetime)
+    resulted_surface = Surface.read_triangular(output_path)
+    assert resulted_surface.creation_datetime == creation_datetime
+    with open(output_path, 'rb') as output_file:
+        assert output_file.read().split(b' on ')[1].startswith(b'Mon Dec 31 21:42:00 2018\n')
+
+
 def test_add_vertex():
     surface = Surface()
     assert not surface.vertices