Type preservation (#164)

* Add failing test * Pass test * lint * Add extra tests
SciTools · Jun 15, 2021 · 7f6cde4 · 7f6cde4
1 parent 02ae952
commit 7f6cde4
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 10 deletions.
diff --git a/cf_units/__init__.py b/cf_units/__init__.py
@@ -1738,10 +1738,12 @@ def convert(self, value, other, ctype=FLOAT64, inplace=False):
                     result, self.cftime_unit, self.calendar)
                 result = cftime.date2num(
                     result_datetimes, other.cftime_unit, other.calendar)
-                # Preserve the datatype of the input array if it was float32.
-                if (isinstance(value, np.ndarray) and
-                   value.dtype == np.float32):
-                    result = result.astype(np.float32)
+                convert_type = (
+                    isinstance(value, np.ndarray)
+                    and np.issubsctype(value.dtype, np.floating)
+                )
+                if convert_type:
+                    result = result.astype(value.dtype)
             else:
                 try:
                     ut_converter = _ud.get_converter(self.ut_unit,

diff --git a/cf_units/tests/unit/unit/test_Unit.py b/cf_units/tests/unit/unit/test_Unit.py
@@ -1,4 +1,4 @@
-# (C) British Crown Copyright 2015 - 2020, Met Office
+# (C) British Crown Copyright 2015 - 2021, Met Office
 #
 # This file is part of cf-units.
 #
@@ -75,11 +75,22 @@ def test_gregorian_calendar_conversion_shape(self):
         self.assertEqual(expected.shape, result.shape)
 
     def test_non_gregorian_calendar_conversion_dtype(self):
-        data = np.arange(4, dtype=np.float32)
-        u1 = Unit('hours since 2000-01-01 00:00:00', calendar='360_day')
-        u2 = Unit('hours since 2000-01-02 00:00:00', calendar='360_day')
-        result = u1.convert(data, u2)
-        self.assertEqual(result.dtype, np.float32)
+        for start_dtype, exp_convert in (
+            (np.float32, True),
+            (np.float64, True),
+            (np.int32, False),
+            (np.int64, False),
+            (np.int, False),
+        ):
+            data = np.arange(4, dtype=start_dtype)
+            u1 = Unit('hours since 2000-01-01 00:00:00', calendar='360_day')
+            u2 = Unit('hours since 2000-01-02 00:00:00', calendar='360_day')
+            result = u1.convert(data, u2)
+
+            if exp_convert:
+                self.assertEqual(result.dtype, start_dtype)
+            else:
+                self.assertEqual(result.dtype, np.int64)
 
 
 class Test_convert__endianness_time(unittest.TestCase):