#137 Add get_field_type

Author: astropenguin

tests/test_typing.py

@@ -8,7 +8,16 @@
 
 
 # submodules
-from xarray_dataclasses.typing import ArrayLike, get_dims, get_dtype
+from xarray_dataclasses.typing import (
+    ArrayLike,
+    Attr,
+    Coord,
+    Data,
+    Name,
+    get_dims,
+    get_dtype,
+    get_field_type,
+)
 
 
 # type hints
@@ -41,6 +50,13 @@
     (ArrayLike[Any, int], "int"),
 ]
 
+testdata_field_type = [
+    (Attr[Any], "attr"),
+    (Coord[Any, Any], "coord"),
+    (Data[Any, Any], "data"),
+    (Name[Any], "name"),
+]
+
 
 # test functions
 @mark.parametrize("type_, dims", testdata_dims)
@@ -51,3 +67,8 @@ def test_get_dims(type_: Any, dims: Any) -> None:
 @mark.parametrize("type_, dtype", testdata_dtype)
 def test_get_dtype(type_: Any, dtype: Any) -> None:
     assert get_dtype(type_) == dtype
+
+
+@mark.parametrize("type_, field_type", testdata_field_type)
+def test_get_field_type(type_: Any, field_type: Any) -> None:
+    assert get_field_type(type_).value == field_type

xarray_dataclasses/typing.py

@@ -19,7 +19,7 @@
 
 # standard library
 from dataclasses import Field
-from enum import auto, Enum
+from enum import Enum
 from typing import (
     Any,
     ClassVar,
@@ -52,22 +52,22 @@
 class FieldType(Enum):
     """Annotation of xarray-related field hints."""
 
-    ATTR = auto()
+    ATTR = "attr"
     """Annotation of attribute field hints."""
 
-    COORD = auto()
+    COORD = "coord"
     """Annotation of coordinate field hints."""
 
-    COORDOF = auto()
+    COORDOF = "coordof"
     """Annotation of coordinate field hints."""
 
-    DATA = auto()
+    DATA = "data"
     """Annotation of data (variable) field hints."""
 
-    DATAOF = auto()
+    DATAOF = "dataof"
     """Annotation of data (variable) field hints."""
 
-    NAME = auto()
+    NAME = "name"
     """Annotation of name field hints."""
 
     def annotates(self, hint: Any) -> bool:
@@ -318,6 +318,23 @@ def get_dtype(type_: Any) -> Dtype:
     raise ValueError(f"Could not convert {type_!r} to dtype.")
 
 
+def get_field_type(type_: Any) -> FieldType:
+    """Parse a type and return a field type if it exists."""
+    if FieldType.ATTR.annotates(type_):
+        return FieldType.ATTR
+
+    if FieldType.COORD.annotates(type_):
+        return FieldType.COORD
+
+    if FieldType.DATA.annotates(type_):
+        return FieldType.DATA
+
+    if FieldType.NAME.annotates(type_):
+        return FieldType.NAME
+
+    raise TypeError(f"Could not find any field type in {type_!r}.")
+
+
 def get_inner(hint: Any, *indexes: int) -> Any:
     """Return an inner type hint by indexes."""
     if not indexes: