TYP: Fix chainmap typing for mypy 0.740+ (#30680)

Author: xhochy

pandas/compat/chainmap.py

@@ -1,23 +1,33 @@
-from collections import ChainMap
+from typing import ChainMap, MutableMapping, TypeVar, cast
 
+_KT = TypeVar("_KT")
+_VT = TypeVar("_VT")
 
-class DeepChainMap(ChainMap):
- def __setitem__(self, key, value):
+
+class DeepChainMap(ChainMap[_KT, _VT]):
+ """Variant of ChainMap that allows direct updates to inner scopes.
+
+ Only works when all passed mapping are mutable.
+ """
+
+ def __setitem__(self, key: _KT, value: _VT) -> None:
  for mapping in self.maps:
- if key in mapping:
- mapping[key] = value
+ mutable_mapping = cast(MutableMapping[_KT, _VT], mapping)
+ if key in mutable_mapping:
+ mutable_mapping[key] = value
  return
- self.maps[0][key] = value
+ cast(MutableMapping[_KT, _VT], self.maps[0])[key] = value
 
- def __delitem__(self, key):
+ def __delitem__(self, key: _KT) -> None:
  """
  Raises
  ------
  KeyError
  If `key` doesn't exist.
  """
  for mapping in self.maps:
+ mutable_mapping = cast(MutableMapping[_KT, _VT], mapping)
  if key in mapping:
- del mapping[key]
+ del mutable_mapping[key]
  return
  raise KeyError(key)

pandas/core/computation/pytables.py

@@ -533,7 +533,7 @@ def __init__(
  self._visitor = None
 
  # capture the environment if needed
- local_dict = DeepChainMap()
+ local_dict: DeepChainMap[Any, Any] = DeepChainMap()
 
  if isinstance(where, PyTablesExpr):
  local_dict = where.env.scope