Minor improvements

README.md

@@ -21,14 +21,14 @@ codemodimportfrom path/to/file.py
 # Overwrite the file
 codemodimportfrom path/to/file.py --write
 
-# Transform ImportFrom statements for modules "foo" and "bar" only
+# Transform ImportFrom statements for modules foo and bar only
 codemodimportfrom path/to/file.py --module foo --module bar
 
-# Allow object imports for "typing.Optional" and "typing.Union"
+# Allow object imports for typing.Optional and typing.Union
 codemodimportfrom path/to/file.py --allow typing.Optional --allow typing.Union
 
 # Allow all object imports from typing
-codemodimportfrom path/to/file.py --allow typing.*
+codemodimportfrom path/to/file.py --allow "typing.*"
 
 # Parse allow list from a .txt file (one line per allowed object import)
 codemodimportfrom path/to/file.py --allow allow.txt

codemodimportfrom/codemodimportfrom.py

@@ -1,3 +1,6 @@
+import sys
+import os
+
 import collections
 import importlib
 import functools
@@ -12,6 +15,7 @@ def transform_importfrom(
     allow: list[str] | None = None,
     transform_module_imports: bool = False,
 ) -> str:
+    sys.path.append(os.getcwd())
     tree = cst.parse_module(code)
     wrapper = cst.metadata.MetadataWrapper(tree)
     tree = wrapper.visit(
@@ -37,6 +41,7 @@ def __init__(
 
     def visit_ImportFrom(self, node: cst.ImportFrom) -> bool | None:
         if node.relative:
+            # TODO handle relative dot imports
             return
         module_name = self._attribute_to_name(node.module)
         if not self.modules or any(
@@ -152,7 +157,7 @@ def _matches_qualified_name_to_leave(self, qualified_name: str) -> bool:
     def _is_module(self, path: str) -> bool:
         try:
             importlib.import_module(path)
-        except ModuleNotFoundError:
+        except Exception:
             return False
         else:
             return True

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "codemodimportfrom"
-version = "0.8.0"
+version = "0.9.0"
 description = ""
 authors = ["Your Name <you@example.com>"]
 readme = "README.md"

tests/test_codemodimportfrom.py

@@ -10,8 +10,8 @@
         #
         ["# foo", "# foo"],
         #
-        ["from . import bar", "from . import bar"],  # TODO handle relative dot imports
-        ["from .bar import baz", "from .bar import baz"],
+        ["from . import bar", "from . import bar"],
+        ["from .foo import bar", "from .foo import bar"],
         #
         [
             """
@@ -94,6 +94,8 @@ def test_rewrites_imports(code, expected_transformed_code):
     assert transformed_code == expected_transformed_code
 
 
+
+
 @pytest.mark.parametrize(
     "code,expected_transformed_code",
     [
@@ -469,3 +471,35 @@ def test_handles_transform_module_imports(code, expected_transformed_code):
     )
 
     assert transformed_code == expected_transformed_code
+
+
+@pytest.mark.parametrize(
+    "code,expected_transformed_code",
+    [
+        [
+            """
+from pydantic import BaseModel, dataclasses
+from pydantic.v1 import BaseModel as V1BaseModel, dataclasses as v1_dataclasses
+BaseModel
+dataclasses
+V1BaseModel
+v1_dataclasses""",
+            """
+from pydantic import BaseModel, dataclasses
+from pydantic.v1 import dataclasses as v1_dataclasses; import pydantic.v1
+BaseModel
+dataclasses
+pydantic.v1.BaseModel
+v1_dataclasses""",
+        ],
+    ],
+)
+def test_handles_dotted_module(code, expected_transformed_code):
+    code = code.strip()
+    expected_transformed_code = expected_transformed_code.strip()
+
+    transformed_code = codemodimportfrom.transform_importfrom(
+        code=code, modules=["pydantic.v1"]
+    )
+
+    assert transformed_code == expected_transformed_code