refactor code for dimension elements

src/cr/cube/dimension.py

@@ -5,7 +5,7 @@
 import copy
 from collections.abc import Sequence
 from datetime import datetime
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union, Callable
 
 import numpy as np
 
@@ -17,6 +17,41 @@
 )
 from cr.cube.util import lazyproperty
 
+DATETIME_FORMATS = {
+    "Y": "%Y",
+    "Q": "%Y-%m",
+    "3M": "%Y-%m",
+    "M": "%Y-%m",
+    "W": "%Y-%m-%d",
+    "D": "%Y-%m-%d",
+    "h": "%Y-%m-%dT%H",
+    "m": "%Y-%m-%dT%H:%M",
+    "s": "%Y-%m-%dT%H:%M:%S",
+    "ms": "%Y-%m-%dT%H:%M:%S.%f",
+    "us": "%Y-%m-%dT%H:%M:%S.%f",
+}
+
+
+def _formatter(dimension_type, typedef, out_format) -> Callable:
+    """Returns a formatting function according to the dimension type."""
+
+    def format(x) -> str:
+        return str(x)
+
+    def format_datetime(x) -> str:
+        try:
+            return datetime.strptime(x, orig_format).strftime(out_format)
+        except ValueError:
+            return str(x)
+
+    if dimension_type != DT.DATETIME:
+        formatter = format
+    else:
+        resolution = typedef["subtype"].get("resolution")
+        orig_format: str = DATETIME_FORMATS.get(resolution) or ""
+        formatter = format_datetime if orig_format and out_format else format
+    return formatter
+
 
 class Dimensions(tuple):
     """Collection containing every dimension defined in cube response."""
@@ -460,20 +495,6 @@ class Elements(tuple):
     Each element is either a category or a subvariable.
     """
 
-    datetime_formats = {
-        "Y": "%Y",
-        "Q": "%Y-%m",
-        "3M": "%Y-%m",
-        "M": "%Y-%m",
-        "W": "%Y-%m-%d",
-        "D": "%Y-%m-%d",
-        "h": "%Y-%m-%dT%H",
-        "m": "%Y-%m-%dT%H:%M",
-        "s": "%Y-%m-%dT%H:%M:%S",
-        "ms": "%Y-%m-%dT%H:%M:%S.%f",
-        "us": "%Y-%m-%dT%H:%M:%S.%f",
-    }
-
     @classmethod
     def from_typedef(
         cls, typedef, dimension_transforms_dict, dimension_type, element_data_format
@@ -509,31 +530,7 @@ def from_typedef(
             xforms = _ElementTransforms(
                 all_xforms.get(element_id, all_xforms.get(str(element_id), {}))
             )
-
-            # This is so dumb that our type checker won't let us just
-            # write `formatter = str`.
-            def format(x) -> str:
-                return str(x)
-
-            formatter = format
-
-            if dimension_type == DT.DATETIME:
-                orig_format: str = (
-                    cls.datetime_formats.get(typedef["subtype"].get("resolution")) or ""
-                )
-                out_format = element_data_format
-                if orig_format and out_format is not None:
-
-                    def format_datetime(x) -> str:
-                        try:
-                            return datetime.strptime(x, orig_format).strftime(
-                                out_format
-                            )
-                        except ValueError:
-                            return str(x)
-
-                    formatter = format_datetime
-
+            formatter = _formatter(dimension_type, typedef, element_data_format)
             element = Element(element_dict, idx, xforms, formatter)
             elements.append(element)