From b7737f29e58df923b1fc6ba28a961a259fa98950 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Fri, 20 Aug 2021 00:15:23 +0200 Subject: [PATCH] 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 --- xarray/core/options.py | 153 +++++++++++++++++++++++++++-------------- 1 file changed, 103 insertions(+), 50 deletions(-) diff --git a/xarray/core/options.py b/xarray/core/options.py index 1d916ff0f7c..524362ce924 100644 --- 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 = ""