Add typing to the OPTIONS dict (#5678)

* Add typing to the OPTIONS dict * Update options.py * Update options.py * Don't use variables as keys * Update options.py * backwards compatibility * Update options.py * Update options.py * Update options.py * get TypedDict from typing_extensions * Update options.py * minor cleanup * Update options.py * Add Colormap as allowed type * linting, keep_attrs can be str or bool * display_expand* can be string or bool * TypedDict is in typing after 3.8 * Try out Literals for better type narrowing * Literal was added python 3.8
pydata · Aug 19, 2021 · b7737f2 · b7737f2
1 parent 378b902
commit b7737f2
Showing 1 changed file with 103 additions and 50 deletions.
diff --git a/xarray/core/options.py b/xarray/core/options.py
@@ -1,38 +1,91 @@
+import sys
 import warnings
 
-ARITHMETIC_JOIN = "arithmetic_join"
-CMAP_DIVERGENT = "cmap_divergent"
-CMAP_SEQUENTIAL = "cmap_sequential"
-DISPLAY_MAX_ROWS = "display_max_rows"
-DISPLAY_STYLE = "display_style"
-DISPLAY_WIDTH = "display_width"
-DISPLAY_EXPAND_ATTRS = "display_expand_attrs"
-DISPLAY_EXPAND_COORDS = "display_expand_coords"
-DISPLAY_EXPAND_DATA_VARS = "display_expand_data_vars"
-DISPLAY_EXPAND_DATA = "display_expand_data"
-ENABLE_CFTIMEINDEX = "enable_cftimeindex"
-FILE_CACHE_MAXSIZE = "file_cache_maxsize"
-KEEP_ATTRS = "keep_attrs"
-WARN_FOR_UNCLOSED_FILES = "warn_for_unclosed_files"
-USE_BOTTLENECK = "use_bottleneck"
-
-
-OPTIONS = {
-    ARITHMETIC_JOIN: "inner",
-    CMAP_DIVERGENT: "RdBu_r",
-    CMAP_SEQUENTIAL: "viridis",
-    DISPLAY_MAX_ROWS: 12,
-    DISPLAY_STYLE: "html",
-    DISPLAY_WIDTH: 80,
-    DISPLAY_EXPAND_ATTRS: "default",
-    DISPLAY_EXPAND_COORDS: "default",
-    DISPLAY_EXPAND_DATA_VARS: "default",
-    DISPLAY_EXPAND_DATA: "default",
-    ENABLE_CFTIMEINDEX: True,
-    FILE_CACHE_MAXSIZE: 128,
-    KEEP_ATTRS: "default",
-    WARN_FOR_UNCLOSED_FILES: False,
-    USE_BOTTLENECK: True,
+# TODO: Remove this check once python 3.7 is not supported:
+if sys.version_info >= (3, 8):
+    from typing import TYPE_CHECKING, Literal, TypedDict, Union
+
+    if TYPE_CHECKING:
+        try:
+            from matplotlib.colors import Colormap
+        except ImportError:
+            Colormap = str
+
+    class T_Options(TypedDict):
+        arithmetic_join: Literal["inner", "outer", "left", "right", "exact"]
+        cmap_divergent: Union[str, "Colormap"]
+        cmap_sequential: Union[str, "Colormap"]
+        display_max_rows: int
+        display_style: Literal["text", "html"]
+        display_width: int
+        display_expand_attrs: Literal["default", True, False]
+        display_expand_coords: Literal["default", True, False]
+        display_expand_data_vars: Literal["default", True, False]
+        display_expand_data: Literal["default", True, False]
+        enable_cftimeindex: bool
+        file_cache_maxsize: int
+        keep_attrs: Literal["default", True, False]
+        warn_for_unclosed_files: bool
+        use_bottleneck: bool
+
+
+else:
+    # See GH5624, this is a convoluted way to allow type-checking to use
+    # `TypedDict` and `Literal` without requiring typing_extensions as a
+    #  required dependency to _run_ the code (it is required to type-check).
+    try:
+        from typing import TYPE_CHECKING, Union
+
+        from typing_extensions import Literal, TypedDict
+
+        if TYPE_CHECKING:
+            try:
+                from matplotlib.colors import Colormap
+            except ImportError:
+                Colormap = str
+
+        class T_Options(TypedDict):
+            arithmetic_join: Literal["inner", "outer", "left", "right", "exact"]
+            cmap_divergent: Union[str, "Colormap"]
+            cmap_sequential: Union[str, "Colormap"]
+            display_max_rows: int
+            display_style: Literal["text", "html"]
+            display_width: int
+            display_expand_attrs: Literal["default", True, False]
+            display_expand_coords: Literal["default", True, False]
+            display_expand_data_vars: Literal["default", True, False]
+            display_expand_data: Literal["default", True, False]
+            enable_cftimeindex: bool
+            file_cache_maxsize: int
+            keep_attrs: Literal["default", True, False]
+            warn_for_unclosed_files: bool
+            use_bottleneck: bool
+
+    except ImportError:
+        from typing import TYPE_CHECKING, Any, Dict, Hashable
+
+        if TYPE_CHECKING:
+            raise
+        else:
+            T_Options = Dict[Hashable, Any]
+
+
+OPTIONS: T_Options = {
+    "arithmetic_join": "inner",
+    "cmap_divergent": "RdBu_r",
+    "cmap_sequential": "viridis",
+    "display_max_rows": 12,
+    "display_style": "html",
+    "display_width": 80,
+    "display_expand_attrs": "default",
+    "display_expand_coords": "default",
+    "display_expand_data_vars": "default",
+    "display_expand_data": "default",
+    "enable_cftimeindex": True,
+    "file_cache_maxsize": 128,
+    "keep_attrs": "default",
+    "warn_for_unclosed_files": False,
+    "use_bottleneck": True,
 }
 
 _JOIN_OPTIONS = frozenset(["inner", "outer", "left", "right", "exact"])
@@ -44,19 +97,19 @@ def _positive_integer(value):
 
 
 _VALIDATORS = {
-    ARITHMETIC_JOIN: _JOIN_OPTIONS.__contains__,
-    DISPLAY_MAX_ROWS: _positive_integer,
-    DISPLAY_STYLE: _DISPLAY_OPTIONS.__contains__,
-    DISPLAY_WIDTH: _positive_integer,
-    DISPLAY_EXPAND_ATTRS: lambda choice: choice in [True, False, "default"],
-    DISPLAY_EXPAND_COORDS: lambda choice: choice in [True, False, "default"],
-    DISPLAY_EXPAND_DATA_VARS: lambda choice: choice in [True, False, "default"],
-    DISPLAY_EXPAND_DATA: lambda choice: choice in [True, False, "default"],
-    ENABLE_CFTIMEINDEX: lambda value: isinstance(value, bool),
-    FILE_CACHE_MAXSIZE: _positive_integer,
-    KEEP_ATTRS: lambda choice: choice in [True, False, "default"],
-    WARN_FOR_UNCLOSED_FILES: lambda value: isinstance(value, bool),
-    USE_BOTTLENECK: lambda value: isinstance(value, bool),
+    "arithmetic_join": _JOIN_OPTIONS.__contains__,
+    "display_max_rows": _positive_integer,
+    "display_style": _DISPLAY_OPTIONS.__contains__,
+    "display_width": _positive_integer,
+    "display_expand_attrs": lambda choice: choice in [True, False, "default"],
+    "display_expand_coords": lambda choice: choice in [True, False, "default"],
+    "display_expand_data_vars": lambda choice: choice in [True, False, "default"],
+    "display_expand_data": lambda choice: choice in [True, False, "default"],
+    "enable_cftimeindex": lambda value: isinstance(value, bool),
+    "file_cache_maxsize": _positive_integer,
+    "keep_attrs": lambda choice: choice in [True, False, "default"],
+    "warn_for_unclosed_files": lambda value: isinstance(value, bool),
+    "use_bottleneck": lambda value: isinstance(value, bool),
 }
 
 
@@ -75,8 +128,8 @@ def _warn_on_setting_enable_cftimeindex(enable_cftimeindex):
 
 
 _SETTERS = {
-    ENABLE_CFTIMEINDEX: _warn_on_setting_enable_cftimeindex,
-    FILE_CACHE_MAXSIZE: _set_file_cache_maxsize,
+    "enable_cftimeindex": _warn_on_setting_enable_cftimeindex,
+    "file_cache_maxsize": _set_file_cache_maxsize,
 }
 
 
@@ -175,9 +228,9 @@ def __init__(self, **kwargs):
                     f"argument name {k!r} is not in the set of valid options {set(OPTIONS)!r}"
                 )
             if k in _VALIDATORS and not _VALIDATORS[k](v):
-                if k == ARITHMETIC_JOIN:
+                if k == "arithmetic_join":
                     expected = f"Expected one of {_JOIN_OPTIONS!r}"
-                elif k == DISPLAY_STYLE:
+                elif k == "display_style":
                     expected = f"Expected one of {_DISPLAY_OPTIONS!r}"
                 else:
                     expected = ""