refactor so UnytScalar defers to a dtype for the unit

Author: ngoldbaum

unytdtype/tests/test_unytdtype.py

@@ -13,6 +13,18 @@ def test_dtype_creation():
     assert dtype == dtype2
 
 
+def test_scalar_creation():
+    dtype = UnytDType("m")
+    unit = unyt.Unit("m")
+    unit_s = "m"
+
+    s_1 = UnytScalar(1, dtype)
+    s_2 = UnytScalar(1, unit)
+    s_3 = UnytScalar(1, unit_s)
+
+    assert s_1 == s_2 == s_3
+
+
 def test_creation_from_zeros():
     dtype = UnytDType("m")
     arr = np.zeros(3, dtype=dtype)

unytdtype/unytdtype/scalar.py

@@ -5,16 +5,24 @@
 
 class UnytScalar:
     def __init__(self, value, unit):
+        from . import UnytDType
         self.value = value
-        if isinstance(unit, str):
-            self.unit = Unit(unit)
-        elif isinstance(unit, Unit):
-            self.unit = unit
+        if isinstance(unit, (str, Unit)):
+            self.dtype = UnytDType(unit)
+        elif isinstance(unit, UnytDType):
+            self.dtype = unit
         else:
             raise RuntimeError
 
+    @property
+    def unit(self):
+        return self.dtype.unit
+
     def __repr__(self):
-        return f"{self.value} {self.unit}"
+        return f"{self.value} {self.dtype.unit}"
 
     def __rmul__(self, other):
-        return UnytScalar(self.value * other, self.unit)
+        return UnytScalar(self.value * other, self.dtype.unit)
+
+    def __eq__(self, other):
+        return self.value == other.value and self.dtype == other.dtype

unytdtype/unytdtype/src/dtype.c

@@ -43,7 +43,12 @@ get_unit(PyObject *scalar, UnytDTypeObject *descr)
         return descr->unit;
     }
 
-    PyObject *unit = PyObject_GetAttrString(scalar, "unit");
+    PyObject *dtype = PyObject_GetAttrString(scalar, "dtype");
+    if (dtype == NULL) {
+        return NULL;
+    }
+    PyObject *unit = PyObject_GetAttrString(dtype, "unit");
+    Py_DECREF(dtype);
     if (unit == NULL) {
         return NULL;
     }