replace bool with bool_t in pandas/core/generic

Author: MarcoGorelli

.pre-commit-config.yaml

@@ -193,6 +193,11 @@ repos:
         language: python
         types: [rst]
         files: ^doc/source/(development|reference)/
+    -   id: no-bool-in-core-generic
+        name: Use bool_t instead of bool in pandas/core/generic.py
+        entry: python scripts/no_bool_in_generic.py
+        language: python
+        files: ^pandas/core/generic\.py$
 -   repo: https://github.com/asottile/yesqa
     rev: v1.2.2
     hooks:

pandas/core/generic.py

@@ -234,7 +234,7 @@ class NDFrame(PandasObject, SelectionMixin, indexing.IndexingMixin):
     def __init__(
         self,
         data: Manager,
-        copy: bool = False,
+        copy: bool_t = False,
         attrs: Optional[Mapping[Optional[Hashable], Any]] = None,
     ):
         # copy kwarg is retained for mypy compat, is not used
@@ -251,7 +251,7 @@ def __init__(
 
     @classmethod
     def _init_mgr(
-        cls, mgr, axes, dtype: Optional[Dtype] = None, copy: bool = False
+        cls, mgr, axes, dtype: Optional[Dtype] = None, copy: bool_t = False
     ) -> Manager:
         """ passed a manager and a axes dict """
         for a, axe in axes.items():
@@ -344,8 +344,8 @@ def flags(self) -> Flags:
     def set_flags(
         self: FrameOrSeries,
         *,
-        copy: bool = False,
-        allows_duplicate_labels: Optional[bool] = None,
+        copy: bool_t = False,
+        allows_duplicate_labels: Optional[bool_t] = None,
     ) -> FrameOrSeries:
         """
         Return a new object with updated flags.
@@ -434,7 +434,7 @@ def _data(self):
     _stat_axis_name = "index"
     _AXIS_ORDERS: List[str]
     _AXIS_TO_AXIS_NUMBER: Dict[Axis, int] = {0: 0, "index": 0, "rows": 0}
-    _AXIS_REVERSED: bool
+    _AXIS_REVERSED: bool_t
     _info_axis_number: int
     _info_axis_name: str
     _AXIS_LEN: int
@@ -461,7 +461,7 @@ def _construct_axes_dict(self, axes=None, **kwargs):
     @final
     @classmethod
     def _construct_axes_from_arguments(
-        cls, args, kwargs, require_all: bool = False, sentinel=None
+        cls, args, kwargs, require_all: bool_t = False, sentinel=None
     ):
         """
         Construct and returns axes if supplied in args/kwargs.
@@ -662,7 +662,7 @@ def _obj_with_exclusions(self: FrameOrSeries) -> FrameOrSeries:
         """ internal compat with SelectionMixin """
         return self
 
-    def set_axis(self, labels, axis: Axis = 0, inplace: bool = False):
+    def set_axis(self, labels, axis: Axis = 0, inplace: bool_t = False):
         """
         Assign desired index to given axis.
 
@@ -693,7 +693,7 @@ def set_axis(self, labels, axis: Axis = 0, inplace: bool = False):
         return self._set_axis_nocheck(labels, axis, inplace)
 
     @final
-    def _set_axis_nocheck(self, labels, axis: Axis, inplace: bool):
+    def _set_axis_nocheck(self, labels, axis: Axis, inplace: bool_t):
         # NDFrame.rename with inplace=False calls set_axis(inplace=True) on a copy.
         if inplace:
             setattr(self, self._get_axis_name(axis), labels)
@@ -932,8 +932,8 @@ def rename(
         index: Optional[Renamer] = None,
         columns: Optional[Renamer] = None,
         axis: Optional[Axis] = None,
-        copy: bool = True,
-        inplace: bool = False,
+        copy: bool_t = True,
+        inplace: bool_t = False,
         level: Optional[Level] = None,
         errors: str = "ignore",
     ) -> Optional[FrameOrSeries]:
@@ -1339,13 +1339,13 @@ def _set_axis_name(self, name, axis=0, inplace=False):
     # Comparison Methods
 
     @final
-    def _indexed_same(self, other) -> bool:
+    def _indexed_same(self, other) -> bool_t:
         return all(
             self._get_axis(a).equals(other._get_axis(a)) for a in self._AXIS_ORDERS
         )
 
     @final
-    def equals(self, other: object) -> bool:
+    def equals(self, other: object) -> bool_t:
         """
         Test whether two objects contain the same elements.
 
@@ -5006,15 +5006,15 @@ def filter(
             return self.reindex(**{name: [r for r in items if r in labels]})
         elif like:
 
-            def f(x) -> bool:
+            def f(x) -> bool_t:
                 assert like is not None  # needed for mypy
                 return like in ensure_str(x)
 
             values = labels.map(f)
             return self.loc(axis=axis)[values]
         elif regex:
 
-            def f(x) -> bool:
+            def f(x) -> bool_t:
                 return matcher.search(ensure_str(x)) is not None
 
             matcher = re.compile(regex)

scripts/no_bool_in_generic.py

@@ -0,0 +1,100 @@
+"""
+Check that pandas/core/generic.py doesn't use bool as a type annotation.
+
+There is already the method `bool`, so the alias `bool_t` should be used instead.
+
+This is meant to be run as a pre-commit hook - to run it manually, you can do:
+
+    pre-commit run no-bool-in-core-generic --all-files
+
+To automatically fixup pandas/core/generic.py, you can pass `--replace`,
+though note that you will also need the additional dependency `tokenize-rt`
+(which is left out from the pre-commit hook so that it uses the same virtualenv
+as the other local ones).
+"""
+
+import argparse
+import ast
+from typing import (
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+)
+
+ERROR_MESSAGE = (
+    "Found annotation 'bool' at line {line}, column {col} - use 'bool_t' instead"
+)
+Offset = Tuple[int, int]
+
+
+class Visitor(ast.NodeVisitor):
+    def __init__(self) -> None:
+        self.to_replace: Set[Offset] = set()
+
+    def generic_visit(self, node):
+        if hasattr(node, "annotation"):
+            if isinstance(node.annotation, ast.Name):
+                if node.annotation.id == "bool":
+                    self.to_replace.add(
+                        (node.annotation.lineno, node.annotation.col_offset)
+                    )
+        elif hasattr(node, "returns"):
+            if isinstance(node.returns, ast.Name):
+                if node.returns.id == "bool":
+                    self.to_replace.add((node.returns.lineno, node.returns.col_offset))
+        super().generic_visit(node)
+
+
+def replace_bool_with_bool_t(visitor: Visitor, content: str) -> str:
+    from tokenize_rt import (
+        reversed_enumerate,
+        src_to_tokens,
+        tokens_to_src,
+    )
+
+    tokens = src_to_tokens(content)
+    for n, i in reversed_enumerate(tokens):
+        if i.offset in visitor.to_replace:
+            tokens[n] = i._replace(src="bool_t")
+
+    new_src: str = tokens_to_src(tokens)
+    return new_src
+
+
+def check_for_bool_in_generic(content: str, *, replace: bool) -> Optional[str]:
+    tree = ast.parse(content)
+
+    visitor = Visitor()
+    visitor.visit(tree)
+
+    if not visitor.to_replace:
+        # Nothing to replace.
+        return content
+
+    if not replace:
+        line, col = visitor.to_replace.pop()
+        msg = ERROR_MESSAGE.format(line=line, col=col)
+        raise RuntimeError(msg)
+
+    return replace_bool_with_bool_t(visitor, content)
+
+
+def main(argv: Optional[Sequence[str]] = None) -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("paths", nargs="*")
+    parser.add_argument("--replace", action="store_true")
+    args = parser.parse_args(argv)
+
+    for path in args.paths:
+        with open(path, encoding="utf-8") as fd:
+            content = fd.read()
+        new_content = check_for_bool_in_generic(content, replace=args.replace)
+        if not args.replace or new_content is None:
+            continue
+        with open(path, "w", encoding="utf-8") as fd:
+            fd.write(new_content)
+
+
+if __name__ == "__main__":
+    main()

scripts/tests/test_no_bool_in_generic.py

@@ -0,0 +1,33 @@
+import pytest
+
+from scripts.no_bool_in_generic import check_for_bool_in_generic
+
+BAD_FILE = "def foo(a: bool) -> bool:\n    ..."
+GOOD_FILE = "def foo(a: bool_t) -> bool_t:\n    ..."
+
+
+def test_bad_file():
+    content = BAD_FILE
+    msg = r"Found annotation 'bool' at line 1, column 11 - use 'bool_t' instead"
+    with pytest.raises(RuntimeError, match=msg):
+        check_for_bool_in_generic(content, replace=False)
+
+
+def test_good_file():
+    # should not raise
+    content = GOOD_FILE
+    check_for_bool_in_generic(content, replace=False)
+
+
+def test_bad_file_with_replace():
+    content = BAD_FILE
+    result = check_for_bool_in_generic(content, replace=True)
+    expected = GOOD_FILE
+    assert result == expected
+
+
+def test_good_file_with_replace():
+    content = GOOD_FILE
+    result = check_for_bool_in_generic(content, replace=True)
+    expected = content
+    assert result == expected