Make pyfiles take a pathlib.Path rather than a str

Author: adamtheturtle

pip_check_reqs/common.py

@@ -124,15 +124,14 @@ def finalise(self) -> Dict[str, FoundModule]:
         return result
 
 
-def pyfiles(root: str) -> Generator[str, None, None]:
-    root_path = Path(root)
-    if root_path.is_file():
-        if root_path.suffix == ".py":
-            yield str(root_path.absolute())
+def pyfiles(root: Path) -> Generator[str, None, None]:
+    if root.is_file():
+        if root.suffix == ".py":
+            yield str(root.absolute())
         else:
-            raise ValueError(f"{root_path} is not a python file or directory")
-    elif root_path.is_dir():
-        for item in root_path.rglob("*.py"):
+            raise ValueError(f"{root} is not a python file or directory")
+    elif root.is_dir():
+        for item in root.rglob("*.py"):
             yield str(item.absolute())
 
 
@@ -143,7 +142,7 @@ def find_imported_modules(
 ) -> Dict[str, FoundModule]:
     vis = _ImportVisitor(ignore_modules_function=ignore_modules_function)
     for path in paths:
-        for filename in pyfiles(path):
+        for filename in pyfiles(Path(path)):
             if ignore_files_function(filename):
                 log.info("ignoring: %s", os.path.relpath(filename))
                 continue

tests/test_common.py

@@ -66,15 +66,15 @@ def test_import_visitor(stmt: str, result: List[str]) -> None:
 def test_pyfiles_file(tmp_path: Path) -> None:
     python_file = tmp_path / "example.py"
     python_file.touch()
-    assert list(common.pyfiles(root=str(python_file))) == [str(python_file)]
+    assert list(common.pyfiles(root=python_file)) == [str(python_file)]
 
 
 def test_pyfiles_file_no_dice(tmp_path: Path) -> None:
     not_python_file = tmp_path / "example"
     not_python_file.touch()
 
     with pytest.raises(ValueError):
-        list(common.pyfiles(root=str(not_python_file)))
+        list(common.pyfiles(root=not_python_file))
 
 
 def test_pyfiles_package(tmp_path: Path) -> None:
@@ -88,7 +88,7 @@ def test_pyfiles_package(tmp_path: Path) -> None:
 
     not_python_file.touch()
 
-    assert list(common.pyfiles(root=str(tmp_path))) == [
+    assert list(common.pyfiles(root=tmp_path)) == [
         str(python_file),
         str(nested_python_file),
     ]