#7 Merge pull request from astropenguin/astropenguin/issue5

Add operator module

CITATION.cff

@@ -3,10 +3,10 @@ message: "If you use this software, please cite it as below."
 
 title: "xarray-units"
 abstract: "xarray extension for handling units"
-version: 0.1.0
-date-released: 2023-12-11
+version: 0.2.0
+date-released: 2023-12-18
 license: "MIT"
-doi: ""
+doi: "10.5281/zenodo.10354517"
 url: "https://github.com/astropenguin/xarray-units/"
 authors:
   - given-names: "Akio"

README.md

@@ -1,2 +1,15 @@
 # xarray-units
+
+[![Release](https://img.shields.io/pypi/v/xarray-units?label=Release&color=cornflowerblue&style=flat-square)](https://pypi.org/project/xarray-units/)
+[![Python](https://img.shields.io/pypi/pyversions/xarray-units?label=Python&color=cornflowerblue&style=flat-square)](https://pypi.org/project/xarray-units/)
+[![Downloads](https://img.shields.io/pypi/dm/xarray-units?label=Downloads&color=cornflowerblue&style=flat-square)](https://pepy.tech/project/xarray-units)
+[![DOI](https://img.shields.io/badge/DOI-10.5281/zenodo.10354517-cornflowerblue?style=flat-square)](https://doi.org/10.5281/zenodo.10354517)
+[![Tests](https://img.shields.io/github/actions/workflow/status/astropenguin/xarray-units/tests.yaml?label=Tests&style=flat-square)](https://github.com/astropenguin/xarray-units/actions)
+
 xarray extension for handling units
+
+## Installation
+
+```shell
+pip install xarray-units==0.2.0
+```

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "xarray-units"
-version = "0.1.0"
+version = "0.2.0"
 description = "xarray extension for handling units"
 authors = ["Akio Taniguchi <[email protected]>"]
 documentation = "https://astropenguin.github.io/xarray-units/"

tests/test_operator.py

@@ -0,0 +1,121 @@
+# standard library
+from typing import Any
+
+
+# dependencies
+from astropy.units import Quantity
+from pytest import mark, raises
+from xarray import DataArray
+from xarray.testing import assert_identical  # type: ignore
+from xarray_units import operator as opr
+from xarray_units.operator import Operator, take
+from xarray_units.quantity import set
+from xarray_units.utils import UnitsApplicationError
+
+
+# test data
+km = set(DataArray([1, 2, 3]), "km")
+mm = set(DataArray([1, 2, 3]) * 1e6, "mm")
+sc_1 = 2
+sc_2 = Quantity(2000, "m")
+
+
+data_take: list[tuple[DataArray, Operator, Any, Any]] = [
+    (km, "mul", sc_1, set(DataArray([2, 4, 6]), "km")),
+    (km, "mul", sc_2, set(DataArray([2, 4, 6]) * 1e3, "km m")),
+    (km, "mul", mm, set(DataArray([1, 4, 9]) * 1e6, "km mm")),
+    (mm, "mul", km, set(DataArray([1, 4, 9]) * 1e6, "km mm")),
+    #
+    (km, "pow", sc_1, set(DataArray([1, 4, 9]), "km2")),
+    (km, "pow", sc_2, UnitsApplicationError),
+    (km, "pow", mm, UnitsApplicationError),
+    (mm, "pow", km, UnitsApplicationError),
+    #
+    (km, "matmul", sc_1, UnitsApplicationError),
+    (km, "matmul", sc_2, UnitsApplicationError),
+    (km, "matmul", mm, set(DataArray(14) * 1e6, "km mm")),
+    (mm, "matmul", km, set(DataArray(14) * 1e6, "km mm")),
+    #
+    (km, "truediv", sc_1, set(DataArray([0.5, 1, 1.5]), "km")),
+    (km, "truediv", sc_2, set(DataArray([0.5, 1.0, 1.5]) * 1e-3, "km m-1")),
+    (km, "truediv", mm, set(DataArray([1, 1, 1]) * 1e-6, "km mm-1")),
+    (mm, "truediv", km, set(DataArray([1, 1, 1]) * 1e6, "mm km-1")),
+    #
+    (km, "add", sc_1, UnitsApplicationError),
+    (km, "add", sc_2, set(DataArray([3, 4, 5]), "km")),
+    (km, "add", mm, set(DataArray([2, 4, 6]), "km")),
+    (mm, "add", km, set(DataArray([2, 4, 6]) * 1e6, "mm")),
+    #
+    (km, "sub", sc_1, UnitsApplicationError),
+    (km, "sub", sc_2, set(DataArray([-1, 0, 1]), "km")),
+    (km, "sub", mm, set(DataArray([0, 0, 0]), "km")),
+    (mm, "sub", km, set(DataArray([0, 0, 0]), "mm")),
+    #
+    (km, "floordiv", sc_1, UnitsApplicationError),
+    (km, "floordiv", sc_2, set(DataArray([0, 1, 1]), "1")),
+    (km, "floordiv", mm, set(DataArray([1, 1, 1]), "1")),
+    (mm, "floordiv", km, set(DataArray([1, 1, 1]), "1")),
+    #
+    (km, "mod", sc_1, UnitsApplicationError),
+    (km, "mod", sc_2, set(DataArray([1, 0, 1]), "km")),
+    (km, "mod", mm, set(DataArray([0, 0, 0]), "km")),
+    (mm, "mod", km, set(DataArray([0, 0, 0]), "mm")),
+    #
+    (km, "lt", sc_1, UnitsApplicationError),
+    (km, "lt", sc_2, DataArray([True, False, False])),
+    (km, "lt", mm, DataArray([False, False, False])),
+    (mm, "lt", km, DataArray([False, False, False])),
+    #
+    (km, "le", sc_1, UnitsApplicationError),
+    (km, "le", sc_2, DataArray([True, True, False])),
+    (km, "le", mm, DataArray([True, True, True])),
+    (mm, "le", km, DataArray([True, True, True])),
+    #
+    (km, "eq", sc_1, UnitsApplicationError),
+    (km, "eq", sc_2, DataArray([False, True, False])),
+    (km, "eq", mm, DataArray([True, True, True])),
+    (mm, "eq", km, DataArray([True, True, True])),
+    #
+    (km, "ne", sc_1, UnitsApplicationError),
+    (km, "ne", sc_2, DataArray([True, False, True])),
+    (km, "ne", mm, DataArray([False, False, False])),
+    (mm, "ne", km, DataArray([False, False, False])),
+    #
+    (km, "ge", sc_1, UnitsApplicationError),
+    (km, "ge", sc_2, DataArray([False, True, True])),
+    (km, "ge", mm, DataArray([True, True, True])),
+    (mm, "ge", km, DataArray([True, True, True])),
+    #
+    (km, "gt", sc_1, UnitsApplicationError),
+    (km, "gt", sc_2, DataArray([False, False, True])),
+    (km, "gt", mm, DataArray([False, False, False])),
+    (mm, "gt", km, DataArray([False, False, False])),
+]
+
+
+@mark.parametrize("left, operator, right, expected", data_take)
+def test_take(
+    left: DataArray,
+    operator: Operator,
+    right: Any,
+    expected: Any,
+) -> None:
+    if expected is UnitsApplicationError:
+        with raises(expected):
+            take(left, operator, right)
+    else:
+        assert_identical(take(left, operator, right), expected)
+
+
+@mark.parametrize("left, operator, right, expected", data_take)
+def test_take_alias(
+    left: DataArray,
+    operator: Operator,
+    right: Any,
+    expected: DataArray,
+) -> None:
+    if expected is UnitsApplicationError:
+        with raises(expected):
+            getattr(opr, operator)(left, right)
+    else:
+        assert_identical(getattr(opr, operator)(left, right), expected)

tests/test_methods.py → tests/test_quantity.py

@@ -3,7 +3,7 @@
 from astropy.constants import c  # type: ignore
 from astropy.units import spectral  # type: ignore
 from xarray.testing import assert_identical  # type: ignore
-from xarray_units.methods import apply, decompose, like, set, to
+from xarray_units.quantity import apply, decompose, like, set, to
 
 
 # test datasets

tests/test_utils.py

@@ -0,0 +1,30 @@
+# standard library
+from typing import Any
+
+
+# dependencies
+from astropy.units import Unit
+from pytest import mark, raises
+from xarray import DataArray
+from xarray_units.utils import UnitsNotValidError, units_of
+
+
+# test data
+data_units_of: list[tuple[Any, Any]] = [
+    (1, None),
+    (Unit("m"), None),
+    (1 * Unit("m"), Unit("m")),
+    (DataArray(1), None),
+    (DataArray(1, attrs={"units": "m"}), Unit("m")),
+    (DataArray(1, attrs={"units": "m, s"}), UnitsNotValidError),
+    (DataArray(1, attrs={"units": "spam"}), UnitsNotValidError),
+]
+
+
+@mark.parametrize("obj, expected", data_units_of)
+def test_units_of(obj: Any, expected: Any) -> None:
+    if expected is UnitsNotValidError:
+        with raises(expected):
+            assert units_of(obj)
+    else:
+        assert units_of(obj) == expected

xarray_units/__init__.py

@@ -1,7 +1,8 @@
-__all__ = ["exceptions", "methods"]
-__version__ = "0.1.0"
+__all__ = ["operator", "quantity", "utils"]
+__version__ = "0.2.0"
 
 
 # submodules
-from . import exceptions
-from . import methods
+from . import operator
+from . import quantity
+from . import utils