ANN: types for _create_storer (pandas-dev#29757)

Author: jbrockmendel

pandas/io/pytables.py

@@ -174,9 +174,6 @@ class DuplicateWarning(Warning):
              and is the default for append operations
 """
 
-# map object types
-_TYPE_MAP = {Series: "series", DataFrame: "frame"}
-
 # storer class map
 _STORER_MAP = {
     "series": "SeriesFixed",
@@ -797,9 +794,10 @@ def select_as_coordinates(
         stop  : integer (defaults to None), row number to stop selection
         """
         where = _ensure_term(where, scope_level=1)
-        return self.get_storer(key).read_coordinates(
-            where=where, start=start, stop=stop, **kwargs
-        )
+        tbl = self.get_storer(key)
+        if not isinstance(tbl, Table):
+            raise TypeError("can only read_coordinates with a table")
+        return tbl.read_coordinates(where=where, start=start, stop=stop, **kwargs)
 
     def select_column(self, key: str, column: str, **kwargs):
         """
@@ -820,7 +818,10 @@ def select_column(self, key: str, column: str, **kwargs):
             is part of a data block)
 
         """
-        return self.get_storer(key).read_column(column=column, **kwargs)
+        tbl = self.get_storer(key)
+        if not isinstance(tbl, Table):
+            raise TypeError("can only read_column with a table")
+        return tbl.read_column(column=column, **kwargs)
 
     def select_as_multiple(
         self,
@@ -903,8 +904,12 @@ def select_as_multiple(
             elif t.nrows != nrows:
                 raise ValueError("all tables must have exactly the same nrows!")
 
+        # The isinstance checks here are redundant with the check above,
+        #  but necessary for mypy; see GH#29757
+        _tbls = [x for x in tbls if isinstance(x, Table)]
+
         # axis is the concentration axes
-        axis = list({t.non_index_axes[0][0] for t in tbls})[0]
+        axis = list({t.non_index_axes[0][0] for t in _tbls})[0]
 
         def func(_start, _stop, _where):
 
@@ -1005,9 +1010,9 @@ def remove(self, key: str, where=None, start=None, stop=None):
                 )
 
             # we are actually trying to remove a node (with children)
-            s = self.get_node(key)
-            if s is not None:
-                s._f_remove(recursive=True)
+            node = self.get_node(key)
+            if node is not None:
+                node._f_remove(recursive=True)
                 return None
 
         # remove the node
@@ -1189,7 +1194,7 @@ def create_table_index(self, key: str, **kwargs):
         if s is None:
             return
 
-        if not s.is_table:
+        if not isinstance(s, Table):
             raise TypeError("cannot create table index on a Fixed format store")
         s.create_index(**kwargs)
 
@@ -1278,7 +1283,7 @@ def get_node(self, key: str):
         except _table_mod.exceptions.NoSuchNodeError:  # type: ignore
             return None
 
-    def get_storer(self, key: str):
+    def get_storer(self, key: str) -> Union["GenericFixed", "Table"]:
         """ return the storer object for a key, raise if not in the file """
         group = self.get_node(key)
         if group is None:
@@ -1331,7 +1336,7 @@ def copy(
                         new_store.remove(k)
 
                 data = self.select(k)
-                if s.is_table:
+                if isinstance(s, Table):
 
                     index: Union[bool, list] = False
                     if propindexes:
@@ -1403,13 +1408,16 @@ def _validate_format(self, format: str, kwargs: Dict[str, Any]) -> Dict[str, Any
 
         return kwargs
 
-    def _create_storer(self, group, format=None, value=None, append=False, **kwargs):
+    def _create_storer(
+        self, group, format=None, value=None, **kwargs
+    ) -> Union["GenericFixed", "Table"]:
         """ return a suitable class to operate """
 
         def error(t):
-            raise TypeError(
+            # return instead of raising so mypy can tell where we are raising
+            return TypeError(
                 f"cannot properly create the storer for: [{t}] [group->"
-                f"{group},value->{type(value)},format->{format},append->{append},"
+                f"{group},value->{type(value)},format->{format},"
                 f"kwargs->{kwargs}]"
             )
 
@@ -1421,6 +1429,7 @@ def error(t):
             if value is None:
 
                 _tables()
+                assert _table_mod is not None  # for mypy
                 if getattr(group, "table", None) or isinstance(
                     group, _table_mod.table.Table
                 ):
@@ -1432,11 +1441,11 @@ def error(t):
                         "nor a value are passed"
                     )
             else:
-
+                _TYPE_MAP = {Series: "series", DataFrame: "frame"}
                 try:
                     pt = _TYPE_MAP[type(value)]
                 except KeyError:
-                    error("_TYPE_MAP")
+                    raise error("_TYPE_MAP")
 
                 # we are actually a table
                 if format == "table":
@@ -1447,7 +1456,7 @@ def error(t):
             try:
                 return globals()[_STORER_MAP[pt]](self, group, **kwargs)
             except KeyError:
-                error("_STORER_MAP")
+                raise error("_STORER_MAP")
 
         # existing node (and must be a table)
         if tt is None:
@@ -1488,7 +1497,7 @@ def error(t):
         try:
             return globals()[_TABLE_MAP[tt]](self, group, **kwargs)
         except KeyError:
-            error("_TABLE_MAP")
+            raise error("_TABLE_MAP")
 
     def _write_to_group(
         self,
@@ -1534,9 +1543,7 @@ def _write_to_group(
                     group = self._handle.create_group(path, p)
                 path = new_path
 
-        s = self._create_storer(
-            group, format, value, append=append, encoding=encoding, **kwargs
-        )
+        s = self._create_storer(group, format, value, encoding=encoding, **kwargs)
         if append:
             # raise if we are trying to append to a Fixed format,
             #       or a table that exists (and we are putting)
@@ -1553,7 +1560,7 @@ def _write_to_group(
         # write the object
         s.write(obj=value, append=append, complib=complib, **kwargs)
 
-        if s.is_table and index:
+        if isinstance(s, Table) and index:
             s.create_index(columns=index)
 
     def _read_group(self, group, **kwargs):
@@ -1584,11 +1591,12 @@ class TableIterator:
 
     chunksize: Optional[int]
     store: HDFStore
+    s: Union["GenericFixed", "Table"]
 
     def __init__(
         self,
         store: HDFStore,
-        s,
+        s: Union["GenericFixed", "Table"],
         func,
         where,
         nrows,
@@ -1651,7 +1659,7 @@ def get_result(self, coordinates: bool = False):
 
         #  return the actual iterator
         if self.chunksize is not None:
-            if not self.s.is_table:
+            if not isinstance(self.s, Table):
                 raise TypeError("can only use an iterator or chunksize on a table")
 
             self.coordinates = self.s.read_coordinates(where=self.where)
@@ -1660,6 +1668,8 @@ def get_result(self, coordinates: bool = False):
 
         # if specified read via coordinates (necessary for multiple selections
         if coordinates:
+            if not isinstance(self.s, Table):
+                raise TypeError("can only read_coordinates on a table")
             where = self.s.read_coordinates(
                 where=self.where, start=self.start, stop=self.stop
             )