coordax 0.2.3__tar.gz → 0.2.4__tar.gz

{coordax-0.2.3 → coordax-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: coordax
-Version: 0.2.3
+Version: 0.2.4
 Summary: Coordinate axes for scientific computing in JAX
 Author-email: Google LLC <noreply@google.com>
 License-Expression: Apache-2.0

{coordax-0.2.3 → coordax-0.2.4}/coordax/__init__.py

@@ -53,4 +53,4 @@ from coordax.ndarrays import (
 )
 import coordax.testing  # pylint: disable=unused-import
 
-__version__ = '0.2.3'  # keep sync with pyproject.toml
+__version__ = '0.2.4'  # keep sync with pyproject.toml

{coordax-0.2.3 → coordax-0.2.4}/coordax/coordinate_systems.py

@@ -25,7 +25,7 @@ import dataclasses
 import functools
 import itertools
 import typing
-from typing import Any, Self, TYPE_CHECKING, Type, TypeAlias, TypeVar
+from typing import Any, Self, TYPE_CHECKING, Type, TypeAlias, TypeGuard, TypeVar
 import warnings
 
 from coordax import utils
@@ -35,7 +35,7 @@ import numpy as np
 if TYPE_CHECKING:
   # import only under TYPE_CHECKING to avoid circular dependency
   # pylint: disable=g-bad-import-order
-  from coordax import fields
+  from coordax import fields  # pylint: disable=unused-import
   import xarray
 
 
@@ -84,7 +84,7 @@ class Coordinate(abc.ABC):
     raise NotImplementedError()
 
   @property
-  def fields(self) -> dict[str, fields.Field]:
+  def fields(self) -> dict[str, 'fields.Field']:
     """Optional dict that maps from field names to their values."""
     return {}
 
@@ -110,7 +110,7 @@ class Coordinate(abc.ABC):
 
   def to_xarray(self) -> dict[str, xarray.Variable]:
     """Convert this coordinate into xarray variables."""
-    import xarray
+    import xarray  # pylint: disable=g-import-not-at-top
 
     variables = {}
     dims_set = {dim for dim in self.dims if dim is not None}
@@ -164,6 +164,12 @@ class ArrayKey:
     return hash((self.value.shape, self.value.tobytes()))
 
 
+@utils.export
+def is_coord(obj: Any) -> TypeGuard[Coordinate]:
+  """Returns True if obj is a Coordinate."""
+  return isinstance(obj, Coordinate)
+
+
 @utils.export
 @jax.tree_util.register_static
 @dataclasses.dataclass(frozen=True)
@@ -211,7 +217,7 @@ class SelectedAxis(Coordinate):
     return (self.coordinate.shape[self.axis],)
 
   @property
-  def fields(self) -> dict[str, fields.Field]:
+  def fields(self) -> dict[str, 'fields.Field']:
     """A maps from field names to their values."""
     return self.coordinate.fields
 
@@ -364,7 +370,7 @@ class CartesianProduct(Coordinate):
     return _concat_tuples(c.shape for c in self.coordinates)
 
   @property
-  def fields(self) -> dict[str, fields.Field]:
+  def fields(self) -> dict[str, 'fields.Field']:
     """Returns a mapping from field names to their values."""
     return _merge_dicts(c.fields for c in self.coordinates)
 
@@ -479,7 +485,7 @@ class LabeledAxis(Coordinate):
     return self.ticks.shape
 
   @property
-  def fields(self) -> dict[str, fields.Field]:
+  def fields(self) -> dict[str, 'fields.Field']:
     # needs local import to avoid circular dependency
     from coordax import fields  # pylint: disable=g-import-not-at-top
 

{coordax-0.2.3 → coordax-0.2.4}/coordax/fields.py

@@ -538,7 +538,7 @@ class Field:
              [0., 0., 0.]], dtype=float32)
       Dimensions without coordinates: x, y
     """
-    import xarray
+    import xarray  # pylint: disable=g-import-not-at-top
 
     if not all(isinstance(dim, str) for dim in self.dims):
       raise ValueError(

{coordax-0.2.3 → coordax-0.2.4}/coordax.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: coordax
-Version: 0.2.3
+Version: 0.2.4
 Summary: Coordinate axes for scientific computing in JAX
 Author-email: Google LLC <noreply@google.com>
 License-Expression: Apache-2.0

{coordax-0.2.3 → coordax-0.2.4}/pyproject.toml

@@ -7,7 +7,7 @@ packages = ["coordax"]
 
 [project]
 name = "coordax"
-version = "0.2.3"  # keep sync with __init__.py
+version = "0.2.4"  # keep sync with __init__.py
 description = "Coordinate axes for scientific computing in JAX"
 authors = [
     {name = "Google LLC", email = "noreply@google.com"},