Merge pull request yt-project#376 from ngoldbaum/setitem-fix

neutrinoceros · Feb 23, 2023 · 5e0cf09 · 5e0cf09
1 parent 2ab7b7b
commit 5e0cf09
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/unyt/array.py b/unyt/array.py
@@ -1729,8 +1729,9 @@ def __getitem__(self, item):
         return ret
 
     def __setitem__(self, item, value):
-        if hasattr(value, "units") and value.units != self.units:
-            value = value.to(self.units)
+        if hasattr(value, "units"):
+            if value.units != self.units and value.units != NULL_UNIT:
+                value = value.to(self.units)
         super().__setitem__(item, value)
 
     def __pow__(self, p, mod=None, /):

diff --git a/unyt/tests/test_unyt_array.py b/unyt/tests/test_unyt_array.py
@@ -2760,3 +2760,7 @@ def test_setitem():
     a[1] = 2
     assert a[1].value == 2
     assert a[1].units == cm
+
+    a[1] = unyt_quantity(2)
+    assert a[1].value == 2
+    assert a[1].units == cm