Add COLS, MAPS and MUT_MAPS to .pyih to be DRY and not lazy

Author: Suor

funcy/colls.pyi

@@ -20,7 +20,7 @@ _T = TypeVar('_T')
 # Non-callable types accepted by the extended function protocol (funcmakers.make_func):
 # int/slice for itemgetter, str/bytes/re.Pattern for regex,
 # Mapping for lookup, Set for membership, None for identity/bool
-# FIX: should include Callable into _Func, it's unied anyway all the time
+# FIX: should include Callable into _Func, it's unied anyway all the time, also name it _XFunc
 _Func = int | slice | str | bytes | re.Pattern[str] | Mapping[Any, Any] | AbstractSet[Any] | None
 _ReResult: TypeAlias = str | tuple[str, ...] | dict[str, str]
 
@@ -43,75 +43,48 @@ def merge_with(f: Callable[[list[Any]], _V], *dicts: Mapping[_K, Any]) -> dict[_
 
 ### Walk / Select
 
-# 1-param collections × XFunc (quadratic expansion)
-# FIX: add abstract collections like MutableMapping, Mapping, Sequence, Set, Iterator, Iterable
-#      drop catch all stuff
-# RES: Sequence, Set, Iterator, Iterable not added - constructors not guaranteed to accept iterables,
-#      and Set/frozenset already covered by concrete expansion above
-# FIX: the above RES is untrue, walk will work with all of the above
-def walk[C: (list, set, frozenset)](f: XFunc[[_T], _V], coll: C[_T]) -> C[_V]: ...
 # dict/MutableMapping/Mapping: typed Callable with pair→pair signature
-# FIX: this also works with collections of pairs, add a test and fix it here
-#      also might need to reorder. Relying on last catch all is not the best solution!
-#      We loose type checking of arguments, i.e. that element types match callable
-def walk(f: Callable[[tuple[_K, _V]], tuple[_K2, _V2]], coll: dict[_K, _V]) -> dict[_K2, _V2]: ...
-def walk(f: Callable[[tuple[_K, _V]], tuple[_K2, _V2]], coll: MutableMapping[_K, _V]) -> dict[_K2, _V2]: ...
-def walk(f: Callable[[tuple[_K, _V]], tuple[_K2, _V2]], coll: Mapping[_K, _V]) -> dict[_K2, _V2]: ...
+# (must come before COLLS — dict is Iterable[K], so Iterable overload would shadow)
+def walk[C: MAPS](f: Callable[[tuple[_K, _V]], tuple[_K2, _V2]], coll: C[_K, _V]) -> C[_K2, _V2]: ...
 # catch-alls for _Func and other collection types
-def walk(f: Callable[..., Any] | _Func, coll: dict[Any, Any]) -> dict[Any, Any]: ...
+def walk[C: MAPS](f: Callable[..., Any] | _Func, coll: C[Any, Any]) -> C[Any, Any]: ...
+# 1-param collections × XFunc (quadratic expansion)
+def walk[C: COLLS](f: XFunc[[_T], _V], coll: C[_T]) -> C[_V]: ...
+# catch-all for other collection types
 def walk(f: Callable[..., Any] | _Func, coll: _T) -> _T: ...
 
 # walk_keys: collection of pairs × XFunc
 # xfunc_skip: slice
-def walk_keys[C: (list, set, frozenset)](f: XFunc[[_K], _K2], coll: C[tuple[_K, _V]]) -> C[tuple[_K2, _V]]: ...
-# walk_keys: dict/MutableMapping/Mapping always returns dict
-# xfunc_skip: slice
-def walk_keys(f: XFunc[[_K], _K2], coll: dict[_K, _V]) -> dict[_K2, _V]: ...
+def walk_keys[C: COLLS](f: XFunc[[_K], _K2], coll: C[tuple[_K, _V]]) -> C[tuple[_K2, _V]]: ...
 # xfunc_skip: slice
-def walk_keys(f: XFunc[[_K], _K2], coll: MutableMapping[_K, _V]) -> dict[_K2, _V]: ...
-# xfunc_skip: slice
-def walk_keys(f: XFunc[[_K], _K2], coll: Mapping[_K, _V]) -> dict[_K2, _V]: ...
+def walk_keys[C: MAPS](f: XFunc[[_K], _K2], coll: C[_K, _V]) -> C[_K2, _V]: ...
 
-# walk_values: collection of pairs × XFunc
-def walk_values[C: (list, set, frozenset)](f: XFunc[[_V], _V2], coll: C[tuple[_K, _V]]) -> C[tuple[_K, _V2]]: ...
-# walk_values: dict/MutableMapping/Mapping always returns dict
-def walk_values(f: XFunc[[_V], _V2], coll: dict[_K, _V]) -> dict[_K, _V2]: ...
-def walk_values(f: XFunc[[_V], _V2], coll: MutableMapping[_K, _V]) -> dict[_K, _V2]: ...
-def walk_values(f: XFunc[[_V], _V2], coll: Mapping[_K, _V]) -> dict[_K, _V2]: ...
+def walk_values[C: COLLS](f: XFunc[[_V], _V2], coll: C[tuple[_K, _V]]) -> C[tuple[_K, _V2]]: ...
+def walk_values[C: MAPS](f: XFunc[[_V], _V2], coll: C[_K, _V]) -> C[_K, _V2]: ...
 
-# 1-param collections × XPred
-def select[C: (list, set, frozenset)](pred: XPred[_T], coll: C[_T]) -> C[_T]: ...
 # dict/Mapping: pred receives (key, value) pairs, only Callable is meaningful
+# (must come before COLLS — dict is Iterable[K], so Iterable overload would shadow these)
 # FIX: should use coll parametrization here
-def select(pred: Callable[[tuple[_K, _V]], Any], coll: dict[_K, _V]) -> dict[_K, _V]: ...
-def select(pred: Callable[[tuple[_K, _V]], Any], coll: MutableMapping[_K, _V]) -> MutableMapping[_K, _V]: ...
-def select(pred: Callable[[tuple[_K, _V]], Any], coll: Mapping[_K, _V]) -> Mapping[_K, _V]: ...
+def select[C: MAPS](pred: Callable[[tuple[_K, _V]], Any], coll: C[_K, _V]) -> C[_K, _V]: ...
+# 1-param collections × XPred
+def select[C: COLLS](pred: XPred[_T], coll: C[_T]) -> C[_T]: ...
 # catch-all for other collection types
 def select(pred: Callable[..., Any] | _Func, coll: _T) -> _T: ...
 # select_keys: collection of pairs × XPred
-def select_keys[C: (list, set, frozenset)](pred: XPred[_K], coll: C[tuple[_K, _V]]) -> C[tuple[_K, _V]]: ...
+def select_keys[C: COLLS](pred: XPred[_K], coll: C[tuple[_K, _V]]) -> C[tuple[_K, _V]]: ...
 # select_keys: dict/MutableMapping/Mapping
-def select_keys[C: (dict, MutableMapping, Mapping)](pred: XPred[_K], coll: C[_K, _V]) -> C[_K, _V]: ...
-
-# FIX0: looks like we are getting lazy everywhere, i.e. here it should be not (list, set, frozenset)
-#       but (list, tuple[T, ...], set, frozenset, Sequence, Iterator, Iterable)
-#       and this is repeating patterns, so let's come up with a systemic decision, we'll add aliases
-#       like ALL_SEQUENCES, ALL_MAPPINGS, ALL_CONTAINERS - these should be global find a good place
-#       for those maybe a separate pyih file like global.pyih or defines.pyih. Then use those aliases
-#       instead of literal listing collection types. Update our pyih translator to support this.
-#
-#       Maybe move other global stuff there too. Could be imports?
+def select_keys[C: MAPS](pred: XPred[_K], coll: C[_K, _V]) -> C[_K, _V]: ...
+
 # select_values: collection of pairs × XPred
-def select_values[C: (list, set, frozenset)](pred: XPred[_V], coll: C[tuple[_K, _V]]) -> C[tuple[_K, _V]]: ...
+def select_values[C: COLLS](pred: XPred[_V], coll: C[tuple[_K, _V]]) -> C[tuple[_K, _V]]: ...
 # select_values: dict/MutableMapping/Mapping
-def select_values[C: (dict, MutableMapping, Mapping)](pred: XPred[_V], coll: C[_K, _V]) -> C[_K, _V]: ...
+def select_values[C: MAPS](pred: XPred[_V], coll: C[_K, _V]) -> C[_K, _V]: ...
 
 def split_keys(pred: XPred[_K], coll: Mapping[_K, _V]) -> tuple[dict[_K, _V], dict[_K, _V]]: ...
 
-def compact[C: (list, set, frozenset)](coll: C[_T]) -> C[_T]: ...
-def compact(coll: dict[_K, _V]) -> dict[_K, _V]: ...
-def compact(coll: MutableMapping[_K, _V]) -> MutableMapping[_K, _V]: ...
-def compact(coll: Mapping[_K, _V]) -> Mapping[_K, _V]: ...
+# (dict/Mapping must come before COLLS — dict is Iterable[K], so Iterable overload would shadow)
+def compact[C: MAPS](coll: C[_K, _V]) -> C[_K, _V]: ...
+def compact[C: COLLS](coll: C[_T]) -> C[_T]: ...
 def compact(coll: _T) -> _T: ...
 
 ### Content tests
@@ -137,10 +110,10 @@ def some(pred: XPred[_T], seq: Iterable[_T]) -> _T | None: ...
 ### Dict utilities
 
 def zipdict(keys: Iterable[_K], vals: Iterable[_V]) -> dict[_K, _V]: ...
-def flip[C: (dict, MutableMapping, Mapping)](mapping: C[_K, _V]) -> C[_V, _K]: ...
-def flip[C: (list, set, frozenset)](coll: C[tuple[_K, _V]]) -> C[tuple[_V, _K]]: ...
-def project[C: (dict, MutableMapping, Mapping)](mapping: C[_K, _V], keys: Iterable[_K]) -> C[_K, _V]: ...
-def omit[C: (dict, MutableMapping, Mapping)](mapping: C[_K, _V], keys: Iterable[_K]) -> C[_K, _V]: ...
+def flip[C: MAPS](mapping: C[_K, _V]) -> C[_V, _K]: ...
+def flip[C: COLLS](coll: C[tuple[_K, _V]]) -> C[tuple[_V, _K]]: ...
+def project[C: MAPS](mapping: C[_K, _V], keys: Iterable[_K]) -> C[_K, _V]: ...
+def omit[C: MAPS](mapping: C[_K, _V], keys: Iterable[_K]) -> C[_K, _V]: ...
 
 def zip_values(*dicts: Mapping[_K, _V]) -> Iterator[tuple[_V, ...]]: ...
 def zip_dicts(*dicts: Mapping[_K, _V]) -> Iterator[tuple[_K, tuple[_V, ...]]]: ...
@@ -149,11 +122,11 @@ def zip_dicts(*dicts: Mapping[_K, _V]) -> Iterator[tuple[_K, tuple[_V, ...]]]: .
 
 def get_in(coll: Mapping[Any, Any] | Sequence[Any], path: Iterable[Any], default: Any = ...) -> Any: ...
 def get_lax(coll: Mapping[Any, Any] | Sequence[Any], path: Iterable[Any], default: Any = ...) -> Any: ...
-def set_in[C: (dict, MutableMapping)](coll: C[_K, _V], path: Sequence[Any], value: Any) -> C[_K, _V]: ...
+def set_in[C: MUT_MAPS](coll: C[_K, _V], path: Sequence[Any], value: Any) -> C[_K, _V]: ...
 def set_in(coll: _T, path: Sequence[Any], value: Any) -> _T: ...
-def update_in[C: (dict, MutableMapping)](coll: C[_K, _V], path: Sequence[Any], update: Callable[[Any], Any], default: Any = ...) -> C[_K, _V]: ...
+def update_in[C: MUT_MAPS](coll: C[_K, _V], path: Sequence[Any], update: Callable[[Any], Any], default: Any = ...) -> C[_K, _V]: ...
 def update_in(coll: _T, path: Sequence[Any], update: Callable[[Any], Any], default: Any = ...) -> _T: ...
-def del_in[C: (dict, MutableMapping)](coll: C[_K, _V], path: Sequence[Any]) -> C[_K, _V]: ...
+def del_in[C: MUT_MAPS](coll: C[_K, _V], path: Sequence[Any]) -> C[_K, _V]: ...
 def del_in(coll: _T, path: Sequence[Any]) -> _T: ...
 def has_path(coll: Mapping[Any, Any] | Sequence[Any], path: Iterable[Any]) -> bool: ...
 

funcy/colls.pyih

@@ -0,0 +1,3 @@
+[mypy]
+# Intentional: more-specific overloads (e.g. dict before Iterable) overlap by design
+disable_error_code = overload-overlap

mypy.ini

@@ -38,7 +38,7 @@ basepython = python3.13
 deps =
     mypy
 commands =
-    python -m mypy.stubtest funcy.types funcy.strings funcy.tree funcy.colls funcy.seqs funcy.funcs funcy.funcolls funcy.flow funcy.calc funcy.debug funcy.objects funcy.decorators --allowlist stubtest_allowlist.txt
+    python -m mypy.stubtest funcy.types funcy.strings funcy.tree funcy.colls funcy.seqs funcy.funcs funcy.funcolls funcy.flow funcy.calc funcy.debug funcy.objects funcy.decorators --allowlist stubtest_allowlist.txt --mypy-config-file {toxinidir}/mypy.ini
 
 [testenv:lint]
 basepython = python3.10

tox.ini

@@ -24,6 +24,13 @@
 
 FUNCY_DIR = Path(__file__).parent / "funcy"
 
+# Collection type aliases for use in .pyih files as [C: ALIAS_NAME]
+COLL_ALIASES = {
+    'COLLS': ['list', 'tuple', 'set', 'frozenset', 'Sequence', 'Iterator', 'Iterable'],
+    'MAPS': ['dict', 'MutableMapping', 'Mapping'],
+    'MUT_MAPS': ['dict', 'MutableMapping'],
+}
+
 # Header for generated files
 GENERATED_HEADER = "# THIS FILE IS AUTOGENERATED by translate_pyih.py from {source}. DO NOT EDIT.\n"
 
@@ -150,10 +157,15 @@ def parse_func_node(node: ast.FunctionDef, xfunc_skip: set) -> dict:
     coll_var = None
     coll_types = None
     for tp in node.type_params:
-        if isinstance(tp, ast.TypeVar) and tp.bound is not None and isinstance(tp.bound, ast.Tuple):
-            coll_var = tp.name
-            coll_types = [annotation_to_str(e) for e in tp.bound.elts]
-            break
+        if isinstance(tp, ast.TypeVar) and tp.bound is not None:
+            if isinstance(tp.bound, ast.Tuple):
+                coll_var = tp.name
+                coll_types = [annotation_to_str(e) for e in tp.bound.elts]
+                break
+            elif isinstance(tp.bound, ast.Name) and tp.bound.id in COLL_ALIASES:
+                coll_var = tp.name
+                coll_types = COLL_ALIASES[tp.bound.id]
+                break
 
     # Build params list, including defaults as part of the type string ("type = default")
     args = node.args
@@ -495,10 +507,44 @@ def expand_collection_types(overload: dict) -> list[dict]:
 
 
 def substitute_coll_type(type_str: str, coll_var: str, ctype: str) -> str:
-    """Replace C[X] or C[X, Y] with concrete_type[X] or concrete_type[X, Y]."""
+    """Replace C[X] or C[X, Y] with concrete_type[X] or concrete_type[X, Y].
+
+    For tuple, C[X] becomes tuple[X, ...] (variable-length homogeneous tuple).
+    """
+    if ctype == 'tuple':
+        # Replace C[...] with tuple[..., ...] using bracket-aware matching
+        result = _substitute_tuple_type(type_str, coll_var)
+        return result
     return re.sub(r'\b' + re.escape(coll_var) + r'\b', ctype, type_str)
 
 
+def _substitute_tuple_type(type_str: str, coll_var: str) -> str:
+    """Replace C[X] with tuple[X, ...] for variable-length tuple semantics."""
+    # FIX: this looks way too complicated, but may go away if we go away from using strings
+    # Find C[ and then match brackets to find the closing ]
+    pattern = re.compile(r'\b' + re.escape(coll_var) + r'\[')
+    result = []
+    pos = 0
+    for m in pattern.finditer(type_str):
+        result.append(type_str[pos:m.start()])
+        # Find matching closing bracket
+        bracket_start = m.end() - 1  # position of '['
+        depth = 1
+        i = bracket_start + 1
+        while i < len(type_str) and depth > 0:
+            if type_str[i] == '[':
+                depth += 1
+            elif type_str[i] == ']':
+                depth -= 1
+            i += 1
+        # Extract inner content (between [ and ])
+        inner = type_str[bracket_start + 1:i - 1]
+        result.append(f'tuple[{inner}, ...]')
+        pos = i
+    result.append(type_str[pos:])
+    return ''.join(result)
+
+
 # ---------------------------------------------------------------------------
 # Output Generation
 # ---------------------------------------------------------------------------

translate_pyih.py

@@ -1,6 +1,18 @@
 from typing import Any, assert_type
 from collections.abc import Iterable, Iterator, Mapping, MutableMapping, Sequence
 
+from funcy import (
+    empty, iteritems, itervalues,
+    join, merge, join_with, merge_with,
+    walk, walk_keys, walk_values, select, select_keys, select_values,
+    split_keys, compact,
+    is_distinct, zipdict, flip, project, omit, zip_values, zip_dicts,
+    where, pluck, pluck_attr, invoke,
+    lwhere, lpluck, lpluck_attr, linvoke,
+    get_in, get_lax, set_in, update_in, del_in, has_path,
+)
+from funcy.colls import all, any, none, one, some  # shadow builtins
+
 # Real abstract-type implementations (not concrete types cast to abstract — checkers see through that)
 class StrIntMapping(Mapping[str, int]):
     def __getitem__(self, k: str) -> int: return 0
@@ -13,17 +25,11 @@ def __setitem__(self, k: str, v: int) -> None: pass
     def __delitem__(self, k: str) -> None: pass
     def __iter__(self) -> Iterator[str]: return iter([])
     def __len__(self) -> int: return 0
-from funcy import (
-    empty, iteritems, itervalues,
-    join, merge, join_with, merge_with,
-    walk, walk_keys, walk_values, select, select_keys, select_values,
-    split_keys, compact,
-    is_distinct, zipdict, flip, project, omit, zip_values, zip_dicts,
-    where, pluck, pluck_attr, invoke,
-    lwhere, lpluck, lpluck_attr, linvoke,
-    get_in, get_lax, set_in, update_in, del_in, has_path,
-)
-from funcy.colls import all, any, none, one, some  # shadow builtins
+
+class IntSequence(Sequence[int]):
+    def __getitem__(self, index: int) -> int: return 0  # type: ignore[override]
+    def __len__(self) -> int: return 0
+
 
 _ReResult = str | tuple[str, ...] | dict[str, str]
 
@@ -122,10 +128,10 @@ def swap_pair(pair: tuple[str, int]) -> tuple[int, str]: return (pair[1], str(pa
 real_mapping = StrIntMapping()
 real_mutable_mapping = StrIntMutableMapping()
 # walk: Mapping with typed pair function returns dict
-reveal_type(walk(swap_pair, real_mapping))  # R: dict[int, str]
+reveal_type(walk(swap_pair, real_mapping))  # R: Mapping[int, str]
 # walk: MutableMapping with typed pair function returns dict
-reveal_type(walk(swap_pair, StrIntMutableMapping()))  # R: dict[int, str]
-reveal_type(walk_keys(str_key_to_int, real_mapping))  # R: dict[int, int]
+reveal_type(walk(swap_pair, StrIntMutableMapping()))  # R: MutableMapping[int, str]
+reveal_type(walk_keys(str_key_to_int, real_mapping))  # R: Mapping[int, int]
 # walk: collection of pairs (list[tuple[K, V]]) — handled by list XFunc overload
 str_int_pairs: list[tuple[str, int]] = [("a", 1), ("b", 2)]
 def transform_pair(p: tuple[str, int]) -> tuple[int, str]: return (p[1], str(p[0]))
@@ -134,8 +140,8 @@ def transform_pair(p: tuple[str, int]) -> tuple[int, str]: return (p[1], str(p[0
 reveal_type(walk(str, str_int_pairs))  # R: list[str]  # XFAIL[ty]: TypeVar inference gives list[tuple[str, int]]
 
 # -- walk_keys: MutableMapping returns dict --
-reveal_type(walk_keys(str_key_to_int, real_mutable_mapping))  # R: dict[int, int]
-reveal_type(walk_keys(None, real_mutable_mapping))  # R: dict[str, int]
+reveal_type(walk_keys(str_key_to_int, real_mutable_mapping))  # R: MutableMapping[int, int]
+reveal_type(walk_keys(None, real_mutable_mapping))  # R: MutableMapping[str, int]
 # -- walk_keys: collection of pairs preserves collection type --
 reveal_type(walk_keys(str_key_to_int, str_int_pairs))  # R: list[tuple[int, int]]
 reveal_type(walk_keys(str.upper, str_int_pairs))  # R: list[tuple[str, int]]
@@ -149,10 +155,10 @@ def transform_pair(p: tuple[str, int]) -> tuple[int, str]: return (p[1], str(p[0
 
 # -- walk_values: always returns dict --
 reveal_type(walk_values(int_to_str, si_dict))  # R: dict[str, str]
-reveal_type(walk_values(int_to_str, real_mapping))  # R: dict[str, str]
+reveal_type(walk_values(int_to_str, real_mapping))  # R: Mapping[str, str]
 # -- walk_values: MutableMapping returns dict --
-reveal_type(walk_values(int_to_str, real_mutable_mapping))  # R: dict[str, str]
-reveal_type(walk_values(None, real_mutable_mapping))  # R: dict[str, int]
+reveal_type(walk_values(int_to_str, real_mutable_mapping))  # R: MutableMapping[str, str]
+reveal_type(walk_values(None, real_mutable_mapping))  # R: MutableMapping[str, int]
 # -- walk_values: collection of pairs preserves collection type --
 reveal_type(walk_values(int_to_str, str_int_pairs))  # R: list[tuple[str, str]]
 reveal_type(walk_values(None, str_int_pairs))  # R: list[tuple[str, int]]
@@ -363,6 +369,33 @@ def add_all(xs: list[int]) -> int: return 0
 reveal_type(walk_keys(None, si_dict))  # R: dict[str, int]
 reveal_type(walk_values(None, si_dict))  # R: dict[str, int]
 
+# -- New collection types: tuple --
+int_tuple: tuple[int, ...] = tuple(int_list)
+reveal_type(walk(int_to_str, int_tuple))  # R: tuple[str, ...]
+reveal_type(walk(None, int_tuple))  # R: tuple[int, ...]
+reveal_type(select(pred_gt0, int_tuple))  # R: tuple[int, ...]
+reveal_type(select(None, int_tuple))  # R: tuple[int, ...]
+reveal_type(compact(int_tuple))  # R: tuple[int, ...]
+str_int_tuple_pairs: tuple[tuple[str, int], ...] = tuple(str_int_pairs)
+reveal_type(walk_keys(str_key_to_int, str_int_tuple_pairs))  # R: tuple[tuple[int, int], ...]
+reveal_type(walk_keys(None, str_int_tuple_pairs))  # R: tuple[tuple[str, int], ...]
+reveal_type(walk_values(int_to_str, str_int_tuple_pairs))  # R: tuple[tuple[str, str], ...]
+reveal_type(select_keys(pred_ne_c, str_int_tuple_pairs))  # R: tuple[tuple[str, int], ...]
+reveal_type(select_values(pred_gt1, str_int_tuple_pairs))  # R: tuple[tuple[str, int], ...]
+reveal_type(flip(str_int_tuple_pairs))  # R: tuple[tuple[int, str], ...]
+
+# -- New collection types: Sequence --
+real_sequence = IntSequence()
+reveal_type(walk(int_to_str, real_sequence))  # R: Sequence[str]
+reveal_type(walk(None, real_sequence))  # R: Sequence[int]
+reveal_type(select(pred_gt0, real_sequence))  # R: Sequence[int]
+reveal_type(compact(real_sequence))  # R: Sequence[int]
+
+# -- New collection types: Iterator --
+int_iter: Iterator[int] = iter([1, 2, 3])
+reveal_type(walk(int_to_str, int_iter))  # R: Iterator[str]  # XFAIL[ty]: Iterator gives Any
+reveal_type(select(pred_gt0, int_iter))  # R: Iterator[int]  # XFAIL[ty]: Iterator gives Any
+
 # -- Should be errors --
 walk(int_to_str, int_list, int_list)  # E: too many arguments
 zipdict(123, [1, 2]) # E: not iterable