Skip to content

Commit b7737f2

Browse files
IllviljanIllviljan
authored
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
1 parent 378b902 commit b7737f2

File tree

1 file changed

+103
-50
lines changed

1 file changed

+103
-50
lines changed

‎xarray/core/options.py‎

Lines changed: 103 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -1,38 +1,91 @@
1+
import sys
12
import warnings
23

3-
ARITHMETIC_JOIN = "arithmetic_join"
4-
CMAP_DIVERGENT = "cmap_divergent"
5-
CMAP_SEQUENTIAL = "cmap_sequential"
6-
DISPLAY_MAX_ROWS = "display_max_rows"
7-
DISPLAY_STYLE = "display_style"
8-
DISPLAY_WIDTH = "display_width"
9-
DISPLAY_EXPAND_ATTRS = "display_expand_attrs"
10-
DISPLAY_EXPAND_COORDS = "display_expand_coords"
11-
DISPLAY_EXPAND_DATA_VARS = "display_expand_data_vars"
12-
DISPLAY_EXPAND_DATA = "display_expand_data"
13-
ENABLE_CFTIMEINDEX = "enable_cftimeindex"
14-
FILE_CACHE_MAXSIZE = "file_cache_maxsize"
15-
KEEP_ATTRS = "keep_attrs"
16-
WARN_FOR_UNCLOSED_FILES = "warn_for_unclosed_files"
17-
USE_BOTTLENECK = "use_bottleneck"
18-
19-
20-
OPTIONS = {
21-
ARITHMETIC_JOIN: "inner",
22-
CMAP_DIVERGENT: "RdBu_r",
23-
CMAP_SEQUENTIAL: "viridis",
24-
DISPLAY_MAX_ROWS: 12,
25-
DISPLAY_STYLE: "html",
26-
DISPLAY_WIDTH: 80,
27-
DISPLAY_EXPAND_ATTRS: "default",
28-
DISPLAY_EXPAND_COORDS: "default",
29-
DISPLAY_EXPAND_DATA_VARS: "default",
30-
DISPLAY_EXPAND_DATA: "default",
31-
ENABLE_CFTIMEINDEX: True,
32-
FILE_CACHE_MAXSIZE: 128,
33-
KEEP_ATTRS: "default",
34-
WARN_FOR_UNCLOSED_FILES: False,
35-
USE_BOTTLENECK: True,
4+
# TODO: Remove this check once python 3.7 is not supported:
5+
if sys.version_info >= (3, 8):
6+
from typing import TYPE_CHECKING, Literal, TypedDict, Union
7+
8+
if TYPE_CHECKING:
9+
try:
10+
from matplotlib.colors import Colormap
11+
except ImportError:
12+
Colormap = str
13+
14+
class T_Options(TypedDict):
15+
arithmetic_join: Literal["inner", "outer", "left", "right", "exact"]
16+
cmap_divergent: Union[str, "Colormap"]
17+
cmap_sequential: Union[str, "Colormap"]
18+
display_max_rows: int
19+
display_style: Literal["text", "html"]
20+
display_width: int
21+
display_expand_attrs: Literal["default", True, False]
22+
display_expand_coords: Literal["default", True, False]
23+
display_expand_data_vars: Literal["default", True, False]
24+
display_expand_data: Literal["default", True, False]
25+
enable_cftimeindex: bool
26+
file_cache_maxsize: int
27+
keep_attrs: Literal["default", True, False]
28+
warn_for_unclosed_files: bool
29+
use_bottleneck: bool
30+
31+
32+
else:
33+
# See GH5624, this is a convoluted way to allow type-checking to use
34+
# `TypedDict` and `Literal` without requiring typing_extensions as a
35+
# required dependency to _run_ the code (it is required to type-check).
36+
try:
37+
from typing import TYPE_CHECKING, Union
38+
39+
from typing_extensions import Literal, TypedDict
40+
41+
if TYPE_CHECKING:
42+
try:
43+
from matplotlib.colors import Colormap
44+
except ImportError:
45+
Colormap = str
46+
47+
class T_Options(TypedDict):
48+
arithmetic_join: Literal["inner", "outer", "left", "right", "exact"]
49+
cmap_divergent: Union[str, "Colormap"]
50+
cmap_sequential: Union[str, "Colormap"]
51+
display_max_rows: int
52+
display_style: Literal["text", "html"]
53+
display_width: int
54+
display_expand_attrs: Literal["default", True, False]
55+
display_expand_coords: Literal["default", True, False]
56+
display_expand_data_vars: Literal["default", True, False]
57+
display_expand_data: Literal["default", True, False]
58+
enable_cftimeindex: bool
59+
file_cache_maxsize: int
60+
keep_attrs: Literal["default", True, False]
61+
warn_for_unclosed_files: bool
62+
use_bottleneck: bool
63+
64+
except ImportError:
65+
from typing import TYPE_CHECKING, Any, Dict, Hashable
66+
67+
if TYPE_CHECKING:
68+
raise
69+
else:
70+
T_Options = Dict[Hashable, Any]
71+
72+
73+
OPTIONS: T_Options = {
74+
"arithmetic_join": "inner",
75+
"cmap_divergent": "RdBu_r",
76+
"cmap_sequential": "viridis",
77+
"display_max_rows": 12,
78+
"display_style": "html",
79+
"display_width": 80,
80+
"display_expand_attrs": "default",
81+
"display_expand_coords": "default",
82+
"display_expand_data_vars": "default",
83+
"display_expand_data": "default",
84+
"enable_cftimeindex": True,
85+
"file_cache_maxsize": 128,
86+
"keep_attrs": "default",
87+
"warn_for_unclosed_files": False,
88+
"use_bottleneck": True,
3689
}
3790

3891
_JOIN_OPTIONS = frozenset(["inner", "outer", "left", "right", "exact"])
@@ -44,19 +97,19 @@ def _positive_integer(value):
4497

4598

4699
_VALIDATORS = {
47-
ARITHMETIC_JOIN: _JOIN_OPTIONS.__contains__,
48-
DISPLAY_MAX_ROWS: _positive_integer,
49-
DISPLAY_STYLE: _DISPLAY_OPTIONS.__contains__,
50-
DISPLAY_WIDTH: _positive_integer,
51-
DISPLAY_EXPAND_ATTRS: lambda choice: choice in [True, False, "default"],
52-
DISPLAY_EXPAND_COORDS: lambda choice: choice in [True, False, "default"],
53-
DISPLAY_EXPAND_DATA_VARS: lambda choice: choice in [True, False, "default"],
54-
DISPLAY_EXPAND_DATA: lambda choice: choice in [True, False, "default"],
55-
ENABLE_CFTIMEINDEX: lambda value: isinstance(value, bool),
56-
FILE_CACHE_MAXSIZE: _positive_integer,
57-
KEEP_ATTRS: lambda choice: choice in [True, False, "default"],
58-
WARN_FOR_UNCLOSED_FILES: lambda value: isinstance(value, bool),
59-
USE_BOTTLENECK: lambda value: isinstance(value, bool),
100+
"arithmetic_join": _JOIN_OPTIONS.__contains__,
101+
"display_max_rows": _positive_integer,
102+
"display_style": _DISPLAY_OPTIONS.__contains__,
103+
"display_width": _positive_integer,
104+
"display_expand_attrs": lambda choice: choice in [True, False, "default"],
105+
"display_expand_coords": lambda choice: choice in [True, False, "default"],
106+
"display_expand_data_vars": lambda choice: choice in [True, False, "default"],
107+
"display_expand_data": lambda choice: choice in [True, False, "default"],
108+
"enable_cftimeindex": lambda value: isinstance(value, bool),
109+
"file_cache_maxsize": _positive_integer,
110+
"keep_attrs": lambda choice: choice in [True, False, "default"],
111+
"warn_for_unclosed_files": lambda value: isinstance(value, bool),
112+
"use_bottleneck": lambda value: isinstance(value, bool),
60113
}
61114

62115

@@ -75,8 +128,8 @@ def _warn_on_setting_enable_cftimeindex(enable_cftimeindex):
75128

76129

77130
_SETTERS = {
78-
ENABLE_CFTIMEINDEX: _warn_on_setting_enable_cftimeindex,
79-
FILE_CACHE_MAXSIZE: _set_file_cache_maxsize,
131+
"enable_cftimeindex": _warn_on_setting_enable_cftimeindex,
132+
"file_cache_maxsize": _set_file_cache_maxsize,
80133
}
81134

82135

@@ -175,9 +228,9 @@ def __init__(self, **kwargs):
175228
f"argument name {k!r} is not in the set of valid options {set(OPTIONS)!r}"
176229
)
177230
if k in _VALIDATORS and not _VALIDATORS[k](v):
178-
if k == ARITHMETIC_JOIN:
231+
if k == "arithmetic_join":
179232
expected = f"Expected one of {_JOIN_OPTIONS!r}"
180-
elif k == DISPLAY_STYLE:
233+
elif k == "display_style":
181234
expected = f"Expected one of {_DISPLAY_OPTIONS!r}"
182235
else:
183236
expected = ""

0 commit comments

Comments
 (0)