def fix_length(length, ds):
registry = ds.unit_registry
if isinstance(length, YTArray):
if registry is not None:
length.units.registry = registry
return length.in_units("code_length")
if isinstance(length, numeric_type):
return YTArray(length, "code_length", registry=registry)
length_valid_tuple = isinstance(length, (list, tuple)) and len(length) == 2
unit_is_string = isinstance(length[1], str)
length_is_number = isinstance(
length[0], numeric_type) and not isinstance(length[0], YTArray)
if length_valid_tuple and unit_is_string and length_is_number:
return YTArray(*length, registry=registry)
else:
raise RuntimeError(f"Length {str(length)} is invalid")
def test_just_one():
# Check that behaviour of this function is consistent before and after refactor
# PR 2893
for unit in ["mm", "cm", "km", "pc", "g", "kg", "M_sun"]:
obj = YTArray([0.0, 1.0], unit)
expected = YTQuantity(obj.flat[0],
obj.units,
registry=obj.units.registry)
jo = just_one(obj)
assert jo == expected