Add type hints

Author: Tatsh

benedict/core/clean.py

@@ -1,4 +1,14 @@
-def _clean_dict(d, strings, collections):
+from collections.abc import MutableMapping, MutableSequence
+from typing import Any, TypeVar
+
+_K = TypeVar("_K")
+_V = TypeVar("_V")
+_T = TypeVar("_T")
+
+
+def _clean_dict(
+    d: MutableMapping[_K, _V], strings: bool, collections: bool
+) -> MutableMapping[_K, _V]:
     keys = list(d.keys())
     for key in keys:
         d[key] = _clean_value(d[key], strings=strings, collections=collections)
@@ -7,41 +17,47 @@ def _clean_dict(d, strings, collections):
     return d
 
 
-def _clean_list(ls, strings, collections):
+def _clean_list(
+    ls: MutableSequence[_T], strings: bool, collections: bool
+) -> MutableSequence[_T]:
     for i in range(len(ls) - 1, -1, -1):
         ls[i] = _clean_value(ls[i], strings=strings, collections=collections)
         if ls[i] is None:
             ls.pop(i)
     return ls
 
 
-def _clean_set(values, strings, collections):
+def _clean_set(values: set[_T], strings: bool, collections: bool) -> set[_T]:
     return {
         value
         for value in values
         if _clean_value(value, strings=strings, collections=collections) is not None
     }
 
 
-def _clean_str(s, strings, collections):
+def _clean_str(s: str, strings: bool, collections: bool) -> str | None:
     return s if s and s.strip() else None
 
 
-def _clean_tuple(values, strings, collections):
+def _clean_tuple(
+    values: tuple[_T, ...], strings: bool, collections: bool
+) -> tuple[_T, ...]:
     return tuple(
         value
         for value in values
         if _clean_value(value, strings=strings, collections=collections) is not None
     )
 
 
-def _clean_value(value, strings, collections):
+def _clean_value(value: Any, strings: bool, collections: bool) -> Any:
     if value is None:
         return value
-    elif isinstance(value, list) and collections:
+    elif isinstance(value, MutableSequence) and collections:
         value = _clean_list(value, strings=strings, collections=collections) or None
-    elif isinstance(value, dict) and collections:
-        value = _clean_dict(value, strings=strings, collections=collections) or None
+    elif isinstance(value, MutableMapping) and collections:
+        value = (
+            _clean_dict(dict(value), strings=strings, collections=collections) or None
+        )
     elif isinstance(value, set) and collections:
         value = _clean_set(value, strings=strings, collections=collections) or None
     elif isinstance(value, str) and strings:
@@ -51,5 +67,5 @@ def _clean_value(value, strings, collections):
     return value
 
 
-def clean(d, strings=True, collections=True):
+def clean(d: Any, strings: bool = True, collections: bool = True) -> Any:
     return _clean_dict(d, strings=strings, collections=collections)

benedict/core/clone.py

@@ -1,8 +1,35 @@
 import copy
+from collections.abc import Mapping, MutableMapping
+from typing import Any, Literal, TypeVar, overload
 
+_K = TypeVar("_K")
+_V = TypeVar("_V")
 
-def clone(obj, empty=False, memo=None):
+
+@overload
+def clone(
+    obj: MutableMapping[_K, _V],
+    empty: Literal[True],
+    memo: dict[int, Any] | None = None,
+) -> MutableMapping[_K, _V]: ...
+
+
+@overload
+def clone(
+    obj: Mapping[_K, _V], empty: Literal[False], memo: dict[int, Any] | None = None
+) -> Mapping[_K, _V]: ...
+
+
+@overload
+def clone(obj: Any, empty: bool = False, memo: dict[int, Any] | None = None) -> Any: ...
+
+
+def clone(
+    obj: Any | MutableMapping[_K, _V] | Mapping[_K, _V],
+    empty: bool = False,
+    memo: dict[int, Any] | None = None,
+) -> MutableMapping[_K, _V] | Mapping[_K, _V] | Any:
     d = copy.deepcopy(obj, memo)
-    if empty:
+    if empty and isinstance(d, MutableMapping):
         d.clear()
     return d

benedict/core/dump.py

@@ -1,7 +1,9 @@
+from typing import Any
+
 from benedict.serializers import JSONSerializer
 
 
-def dump(obj, **kwargs):
+def dump(obj: Any, **kwargs: Any) -> str:
     serializer = JSONSerializer()
     options = {"indent": 4, "sort_keys": True}
     options.update(**kwargs)

benedict/core/filter.py

@@ -1,13 +1,20 @@
-from benedict.core import clone
+from collections.abc import Callable, MutableMapping
+from typing import TypeVar
 
+from benedict.core.clone import clone
 
-def filter(d, predicate):
+_V = TypeVar("_V")
+
+
+def filter(
+    d: MutableMapping[str, _V | None], predicate: Callable[[str, _V | None], bool]
+) -> MutableMapping[str, _V | None]:
     if not callable(predicate):
         raise ValueError("predicate argument must be a callable.")
     new_dict = clone(d, empty=True)
     keys = list(d.keys())
     for key in keys:
-        value = d.get(key, None)
+        value = d.get(key)
         if predicate(key, value):
             new_dict[key] = value
     return new_dict

benedict/core/find.py

@@ -1,5 +1,14 @@
-def find(d, keys, default=None):
+from collections.abc import Iterable, Mapping
+from typing import TypeVar
+
+_K = TypeVar("_K")
+_V = TypeVar("_V")
+
+
+def find(
+    d: Mapping[_K, _V], keys: Iterable[_K], default: _V | None = None
+) -> _V | None:
     for key in keys:
         if key in d:
-            return d.get(key, default)
+            return d[key]
     return default

benedict/core/flatten.py

@@ -1,20 +1,29 @@
-from benedict.core import clone
-from benedict.utils import type_util
+from collections.abc import Mapping, MutableMapping
+from typing import TypeVar
 
+from benedict.core.clone import clone
 
-def _flatten_key(base_key, key, separator):
+_V = TypeVar("_V")
+
+
+def _flatten_key(base_key: str, key: str, separator: str) -> str:
     if base_key and separator:
         return f"{base_key}{separator}{key}"
     return key
 
 
-def _flatten_item(d, base_dict, base_key, separator):
+def _flatten_item(
+    d: Mapping[str, _V],
+    base_dict: MutableMapping[str, _V],
+    base_key: str,
+    separator: str,
+) -> MutableMapping[str, _V]:
     new_dict = base_dict
     keys = list(d.keys())
     for key in keys:
         new_key = _flatten_key(base_key, key, separator)
-        value = d.get(key, None)
-        if type_util.is_dict(value):
+        value = d.get(key, None)  # type: ignore[arg-type]
+        if isinstance(value, Mapping):
             new_value = _flatten_item(
                 value, base_dict=new_dict, base_key=new_key, separator=separator
             )
@@ -26,6 +35,8 @@ def _flatten_item(d, base_dict, base_key, separator):
     return new_dict
 
 
-def flatten(d, separator="_"):
-    new_dict = clone(d, empty=True)
+def flatten(
+    d: MutableMapping[str, _V], separator: str = "_"
+) -> MutableMapping[str, _V]:
+    new_dict: MutableMapping[str, _V] = clone(d, empty=True)
     return _flatten_item(d, base_dict=new_dict, base_key="", separator=separator)

benedict/core/groupby.py

@@ -1,14 +1,22 @@
+from collections.abc import Mapping, MutableSequence, Sequence
+from typing import Any, TypeVar
+
 from benedict.utils import type_util
 
+_K = TypeVar("_K")
+_V = TypeVar("_V", bound=MutableSequence[Any])
+
 
-def groupby(items, key):
+def groupby(
+    items: Sequence[Mapping[_K, _V | None]], key: _K
+) -> dict[_K | None, _V | list[Any]]:
     if not type_util.is_list(items):
         raise ValueError("items should be a list of dicts.")
-    items_grouped = {}
+    items_grouped: dict[_K | None, _V | list[Any]] = {}
     for item in items:
         if not type_util.is_dict(item):
             raise ValueError("item should be a dict.")
-        group = item.get(key)
+        group: _K | None = item.get(key)
         if group not in items_grouped:
             items_grouped[group] = []
         items_grouped[group].append(item.copy())

benedict/core/invert.py

@@ -1,22 +1,35 @@
-from benedict.core import clone
+from collections.abc import MutableMapping, Sequence
+from typing import Any, TypeVar
+
+from benedict.core.clone import clone
 from benedict.utils import type_util
 
+_K = TypeVar("_K")
+_V = TypeVar("_V")
+
 
-def _invert_item(d, key, value, flat):
+def _invert_item(
+    d: MutableMapping[_V, _K | list[_K]], key: _K, value: _V, flat: bool
+) -> None:
     if flat:
-        d.setdefault(value, key)
+        d[value] = key
     else:
-        d.setdefault(value, []).append(key)
+        d[value] = [key]
 
 
-def _invert_list(d, key, value, flat):
+def _invert_list(
+    d: MutableMapping[_V, _K | list[_K]], key: _K, value: Sequence[Any], flat: bool
+) -> None:
     for value_item in value:
         _invert_item(d, key, value_item, flat)
 
 
-def invert(d, flat=False):
-    new_dict = clone(d, empty=True)
-    for key, value in d.items():
+def invert(
+    d: MutableMapping[_K, _V], flat: bool = False
+) -> MutableMapping[_V, _K | list[_K]]:
+    orig_dict = clone(d, empty=True)
+    new_dict: dict[_V, _K | list[_K]] = {}
+    for key, value in orig_dict.items():
         if type_util.is_list_or_tuple(value):
             _invert_list(new_dict, key, value, flat)
         else:

benedict/core/items_sorted.py

@@ -1,10 +1,23 @@
-def _items_sorted_by_item_at_index(d, index, reverse):
+from collections.abc import Mapping
+
+from useful_types import SupportsRichComparisonT
+
+
+def _items_sorted_by_item_at_index(
+    d: Mapping[SupportsRichComparisonT, SupportsRichComparisonT],
+    index: int,
+    reverse: bool,
+) -> list[tuple[SupportsRichComparisonT, SupportsRichComparisonT]]:
     return sorted(d.items(), key=lambda item: item[index], reverse=reverse)
 
 
-def items_sorted_by_keys(d, reverse=False):
+def items_sorted_by_keys(
+    d: Mapping[SupportsRichComparisonT, SupportsRichComparisonT], reverse: bool = False
+) -> list[tuple[SupportsRichComparisonT, SupportsRichComparisonT]]:
     return _items_sorted_by_item_at_index(d, 0, reverse)
 
 
-def items_sorted_by_values(d, reverse=False):
+def items_sorted_by_values(
+    d: Mapping[SupportsRichComparisonT, SupportsRichComparisonT], reverse: bool = False
+) -> list[tuple[SupportsRichComparisonT, SupportsRichComparisonT]]:
     return _items_sorted_by_item_at_index(d, 1, reverse)

benedict/core/keylists.py

@@ -1,7 +1,12 @@
+from collections.abc import Mapping, Sequence
+from typing import Any
+
 from benedict.utils import type_util
 
 
-def _get_keylist_for_dict(d, parent_keys, indexes):
+def _get_keylist_for_dict(
+    d: Mapping[Any, Any], parent_keys: list[Any], indexes: bool
+) -> list[list[Any]]:
     keylist = []
     for key, value in d.items():
         keys = parent_keys + [key]
@@ -10,7 +15,9 @@ def _get_keylist_for_dict(d, parent_keys, indexes):
     return keylist
 
 
-def _get_keylist_for_list(ls, parent_keys, indexes):
+def _get_keylist_for_list(
+    ls: Sequence[Any], parent_keys: list[Any], indexes: bool
+) -> list[list[Any]]:
     keylist = []
     for key, value in enumerate(ls):
         keys = list(parent_keys)
@@ -20,13 +27,17 @@ def _get_keylist_for_list(ls, parent_keys, indexes):
     return keylist
 
 
-def _get_keylist_for_value(value, parent_keys, indexes):
+def _get_keylist_for_value(
+    value: Mapping[Any, Any] | Sequence[Any], parent_keys: list[Any], indexes: bool
+) -> list[list[Any]]:
     if type_util.is_dict(value):
         return _get_keylist_for_dict(value, parent_keys, indexes)
     elif type_util.is_list(value) and indexes:
         return _get_keylist_for_list(value, parent_keys, indexes)
     return []
 
 
-def keylists(d, indexes=False):
+def keylists(
+    d: Mapping[Any, Any] | Sequence[Any], indexes: bool = False
+) -> list[list[Any]]:
     return _get_keylist_for_value(d, [], indexes)