minor improvements to check (#7036)

dagster-io · Mar 14, 2022 · 756527a · 756527a
1 parent a88e658
commit 756527a
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 27 deletions.
diff --git a/python_modules/dagster/dagster/check/__init__.py b/python_modules/dagster/dagster/check/__init__.py
@@ -282,44 +282,54 @@ def opt_two_dim_dict_param(
     return _check_two_dim_mapping_entries(obj, key_type, value_type, mapping_type=dict)
 
 
-def dict_elem(ddict: Dict, key: str) -> Dict:
+def dict_elem(
+    obj: Dict,
+    key: str,
+    key_type: Optional[TypeOrTupleOfTypes] = None,
+    value_type: Optional[TypeOrTupleOfTypes] = None,
+) -> Dict:
     from dagster.utils import frozendict
 
-    dict_param(ddict, "ddict")
+    dict_param(obj, "obj")
     str_param(key, "key")
 
-    if key not in ddict:
-        raise CheckError(f"{key} not present in dictionary {ddict}")
+    if key not in obj:
+        raise CheckError(f"{key} not present in dictionary {obj}")
 
-    value = ddict[key]
+    value = obj[key]
     if not isinstance(value, (frozendict, dict)):
-        raise _element_check_error(key, value, ddict, (frozendict, dict))
-    return value
+        raise _element_check_error(key, value, obj, (frozendict, dict))
+    else:
+        return _check_mapping_entries(value, key_type, value_type, mapping_type=dict)
 
 
-def opt_dict_elem(ddict: Dict, key: str) -> Dict:
+def opt_dict_elem(
+    obj: Dict[str, Any],
+    key: str,
+    key_type: Optional[TypeOrTupleOfTypes] = None,
+    value_type: Optional[TypeOrTupleOfTypes] = None,
+) -> Dict:
     from dagster.utils import frozendict
 
-    dict_param(ddict, "ddict")
+    dict_param(obj, "obj")
     str_param(key, "key")
 
-    value = ddict.get(key)
+    value = obj.get(key)
 
     if value is None:
         return {}
-
-    if not isinstance(value, (frozendict, dict)):
-        raise _element_check_error(key, value, ddict, list)
-
-    return value
+    elif not isinstance(value, (frozendict, dict)):
+        raise _element_check_error(key, value, obj, dict)
+    else:
+        return _check_mapping_entries(value, key_type, value_type, mapping_type=dict)
 
 
 def is_dict(
-    obj: object,
+    obj: Dict[T, U],
     key_type: Optional[TypeOrTupleOfTypes] = None,
     value_type: Optional[TypeOrTupleOfTypes] = None,
     desc: Optional[str] = None,
-) -> Dict:
+) -> Dict[T, U]:
     from dagster.utils import frozendict
 
     if not isinstance(obj, (frozendict, dict)):
@@ -885,8 +895,8 @@ def opt_nullable_sequence_param(
 
 
 def set_param(
-    obj: object, param_name: str, of_type: Optional[TypeOrTupleOfTypes] = None
-) -> AbstractSet:
+    obj: AbstractSet[T], param_name: str, of_type: Optional[TypeOrTupleOfTypes] = None
+) -> AbstractSet[T]:
     if not isinstance(obj, (frozenset, set)):
         raise _param_type_mismatch_exception(obj, (frozenset, set), param_name)
 
@@ -897,27 +907,41 @@ def set_param(
 
 
 def opt_set_param(
-    obj: object, param_name: str, of_type: Optional[TypeOrTupleOfTypes] = None
-) -> AbstractSet:
+    obj: Optional[AbstractSet[T]], param_name: str, of_type: Optional[TypeOrTupleOfTypes] = None
+) -> AbstractSet[T]:
     """Ensures argument obj is a set or None; in the latter case, instantiates an empty set
     and returns it.
 
     If the of_type argument is provided, also ensures that list items conform to the type specified
     by of_type.
     """
-    if obj is not None and not isinstance(obj, (frozenset, set)):
-        raise _param_type_mismatch_exception(obj, (frozenset, set), param_name)
-    if not obj:
+    if obj is None:
         return set()
-    if not of_type:
+    elif obj is not None and not isinstance(obj, (frozenset, set)):
+        raise _param_type_mismatch_exception(obj, (frozenset, set), param_name)
+    elif not of_type:
         return obj
 
     return _check_iterable_items(obj, of_type, "set")
 
 
+@overload
 def opt_nullable_set_param(
-    obj: object, param_name: str, of_type: Optional[TypeOrTupleOfTypes] = None
-) -> Optional[AbstractSet]:
+    obj: None, param_name: str, of_type: Optional[TypeOrTupleOfTypes] = ...
+) -> None:
+    ...
+
+
+@overload
+def opt_nullable_set_param(
+    obj: AbstractSet[T], param_name: str, of_type: Optional[TypeOrTupleOfTypes] = ...
+) -> AbstractSet[T]:
+    ...
+
+
+def opt_nullable_set_param(
+    obj: Optional[AbstractSet[T]], param_name: str, of_type: Optional[TypeOrTupleOfTypes] = None
+) -> Optional[AbstractSet[T]]:
     """Ensures argument obj is a set or None. Returns None if input is None.
     and returns it.
 

diff --git a/python_modules/dagster/dagster_tests/general_tests/check_tests/test_check.py b/python_modules/dagster/dagster_tests/general_tests/check_tests/test_check.py
@@ -431,18 +431,29 @@ def test_dict_elem():
     with pytest.raises(CheckError):
         check.dict_elem(ddict, "nonexistantkey")
 
+    assert check.dict_elem(
+        {"foo": {"str": 1, "str2": "str", 1: "str", 2: "str"}},
+        "foo",
+        key_type=(str, int),
+        value_type=(str, int),
+    )
+
 
 def test_opt_dict_elem():
     dict_value = {"blah": "blahblah"}
     ddict = {"dictkey": dict_value, "stringkey": "A", "nonekey": None}
 
     assert check.opt_dict_elem(ddict, "dictkey") == dict_value
+    assert check.opt_dict_elem(ddict, "dictkey", key_type=str, value_type=str) == dict_value
     assert check.opt_dict_elem(ddict, "nonekey") == {}
     assert check.opt_dict_elem(ddict, "nonexistantkey") == {}
 
     with pytest.raises(CheckError):
         check.opt_dict_elem(ddict, "stringkey")
 
+    with pytest.raises(CheckError):
+        check.opt_dict_elem(ddict, "dictkey", key_type=str, value_type=int)
+
 
 # ########################
 # ##### FLOAT