Improve and document the import hook mechanism (#5)

* Improve import hook

* Add tests

* Update examples

* Update readme

* Remove coding directive from import hook examples

* Make pyright happy

README.md

@@ -33,7 +33,13 @@ Get it via pip:
 pip install python-jsx
 ```
 
-## Minimal example
+## Minimal example (using the `coding` directive)
+
+> [!TIP]
+> There are more examples available in the [examples folder](examples).
+
+There are two supported ways to seamlessly integrate JSX into your codebase.
+One is by registering a custom codec shown here and the other by using a custom import hook shown [below](#minimal-example-using-an-import-hook).
 
 ```python
 # hello.py
@@ -58,9 +64,6 @@ $ python main.py
 <h1>Hello, word!</h1>
 ```
 
-> [!TIP]
-> There are more examples available in the [examples folder](examples).
-
 Each file containing JSX must contain two things:
 
 - `# coding: jsx` directive - This tells Python to let our library parse the
@@ -72,6 +75,42 @@ To run a file containing JSX, the `jsx` codec must be registered first which can
 be done with `from pyjsx import auto_setup`. This must occur before importing
 any other file containing JSX.
 
+## Minimal example (using an import hook)
+
+> [!TIP]
+> There are more examples available in the [examples folder](examples).
+
+```python
+# hello.px
+from pyjsx import jsx
+
+def hello():
+    print(<h1>Hello, world!</h1>)
+```
+
+```python
+# main.py
+from pyjsx import auto_setup
+
+from hello import hello
+
+hello()
+```
+
+```sh
+$ python main.py
+<h1>Hello, word!</h1>
+```
+
+Each file containing JSX must contain two things:
+
+- The file extension must be `.px`
+- `from pyjsx import jsx` import. PyJSX transpiles JSX into `jsx(...)` calls so
+  it must be in scope.
+
+To be able to import `.px`, the import hook must be registered first which can
+be done with `from pyjsx import auto_setup` (same as for the codec version). This must occur before importing any other file containing JSX.
+
 ## Supported grammar
 
 The full [JSX grammar](https://facebook.github.io/jsx/) is supported.

examples/README.md

@@ -3,6 +3,9 @@
 > [!TIP]
 > Run each example with `python main.py`
 
-- table - Shows how you can easily generate an HTML table from data
-- custom - Shows how you can use custom components
-- props - Shows some advanced props usage
+The examples showcase the two supported ways of running JSX in Python.
+Examples with `_codec` show how to use a custom codec. Examples with `_import_hook` show how to use an import hook.
+
+- `table` - Shows how you can easily generate an HTML table from data
+- `custom_components` - Shows how you can use custom components
+- `props` - Shows some advanced props usage

examples/custom_components_import_hook/custom.px

@@ -0,0 +1,20 @@
+from pyjsx import jsx, JSX
+
+
+def Header(children, style=None, **rest) -> JSX:
+    return <h1 style={style}>{children}</h1>
+
+
+def Main(children, **rest) -> JSX:
+    return <main>{children}</main>
+
+
+def App() -> JSX:
+    return (
+        <div>
+            <Header style={{"color": "red"}}>Hello, world!</Header>
+            <Main>
+                <p>This was rendered with PyJSX!</p>
+            </Main>
+        </div>
+    )

examples/custom_components_import_hook/main.py

@@ -0,0 +1,6 @@
+from pyjsx import auto_setup
+
+from custom import App
+
+
+print(App())

examples/props_import_hook/main.py

@@ -0,0 +1,5 @@
+from pyjsx import auto_setup
+
+from props import App
+
+print(App())

examples/props_import_hook/props.px

@@ -0,0 +1,37 @@
+from pyjsx import jsx, JSX
+
+
+def Card(rounded=False, raised=False, image=None, children=None, **rest) -> JSX:
+    style = {
+        "border-radius": "5px" if rounded else 0,
+        "box-shadow": "0 2px 4px rgba(0, 0, 0, 0.1)" if raised else "none",
+    }
+    return (
+        <div style={style}>
+            {image}
+            {children}
+        </div>
+    )
+
+
+def Image(src, alt, **rest) -> JSX:
+    return <img src={src} alt={alt} />
+
+
+def App() -> JSX:
+    return (
+        <div>
+            <Card rounded raised image={<Image src="dog.jpg" alt="A picture of a dog" />}>
+                <h1>Card title</h1>
+                <p>Card content</p>
+            </Card>
+            <Card rounded raised={False} disabled image={<Image src="cat.jpg" alt="A picture of a cat" />}>
+                <h1>Card title</h1>
+                <p>Card content</p>
+            </Card>
+            <Card rounded raised={False}>
+                <h1>Card title</h1>
+                <p>Card content</p>
+            </Card>
+        </div>
+    )

examples/table_import_hook/main.py

@@ -0,0 +1,6 @@
+from pyjsx import auto_setup
+
+from table import make_table
+
+
+print(make_table())

examples/table_import_hook/table.px

@@ -0,0 +1,36 @@
+from pyjsx import jsx, JSX
+
+
+def make_header(names: list[str]) -> JSX:
+    return (
+        <thead>
+            <tr>
+                {<th>{name}</th> for name in names}
+            </tr>
+        </thead>
+    )
+
+
+def make_body(rows: list[list[str]]) -> JSX:
+    return (
+        <tbody>
+            {
+                <tr>
+                    {<td>{cell}</td> for cell in row}
+                </tr>
+                for row in rows
+            }
+        </tbody>
+    )
+
+
+def make_table() -> JSX:
+    columns = ["Name", "Age"]
+    rows = [["Alice", "34"], ["Bob", "56"]]
+
+    return (
+        <table>
+            {make_header(columns)}
+            {make_body(rows)}
+        </table>
+    )

pyjsx/import_hook.py

@@ -23,11 +23,10 @@
 from pyjsx.transpiler import transpile
 
 
-class PyJSXLoader(FileLoader):
-    def __init__(self, name: str):
-        self.name = name
-        self.path = f"{name}.px"
+PYJSX_SUFFIX = ".px"
+
 
+class PyJSXLoader(FileLoader):
     def _compile(self) -> str:
         return transpile(Path(self.path).read_text("utf-8"))
 
@@ -46,14 +45,26 @@ def find_spec(
         path: Sequence[str] | None,
         target: ModuleType | None = None,  # noqa: ARG002
     ) -> ModuleSpec | None:
-        filename = f"{fullname}.px"
-        if not Path(filename).exists():
-            return None
-        if path:
-            msg = "Only top-level imports are supported"
-            raise NotImplementedError(msg)
-        return importlib.util.spec_from_loader(fullname, PyJSXLoader(fullname))
+        if not path:
+            path = sys.path
+
+        for p in path:
+            if spec := self._spec_from_path(fullname, p):
+                return spec
+
+    def _spec_from_path(self, fullname: str, path: str) -> ModuleSpec | None:
+        last_segment = fullname.rsplit(".", maxsplit=1)[-1]
+        full_path = Path(path) / f"{last_segment}{PYJSX_SUFFIX}"
+        if full_path.exists():
+            loader = PyJSXLoader(fullname, str(full_path))
+            return importlib.util.spec_from_loader(fullname, loader)
 
 
 def register_import_hook() -> None:
+    """Register import hook for .px files."""
     sys.meta_path.append(PyJSXFinder())
+
+
+def unregister_import_hook() -> None:
+    """Unregister import hook for .px files."""
+    sys.meta_path = [finder for finder in sys.meta_path if not isinstance(finder, PyJSXFinder)]

tests/test_examples.py

@@ -13,15 +13,18 @@ def run_example(name: str):
 
 
 @pytest.mark.parametrize(
-    "example",
+    ("example", "loader"),
     [
-        "table",
-        "props",
-        "custom",
+        ("table", "codec"),
+        ("table", "import_hook"),
+        ("props", "codec"),
+        ("props", "import_hook"),
+        ("custom_components", "codec"),
+        ("custom_components", "import_hook"),
     ],
 )
-def test_example(request, snapshot, example):
+def test_example(snapshot, example, loader):
     snapshot.snapshot_dir = Path(__file__).parent / "data"
     snapshot.assert_match(
-        run_example(example), f"examples-{request.node.callspec.id}.txt"
+        run_example(f"{example}_{loader}"), f"examples-{example}.txt"
     )

tests/test_import_hook.py

@@ -0,0 +1,30 @@
+from pathlib import Path
+
+import pytest
+
+from pyjsx.import_hook import PyJSXFinder, register_import_hook, unregister_import_hook
+
+
+@pytest.fixture
+def import_hook():
+    register_import_hook()
+    yield
+    unregister_import_hook()
+
+
+def test_finder():
+    finder = PyJSXFinder()
+    path = str(Path(__file__).parent / "test_module")
+    spec = finder.find_spec("main", [path])
+    assert spec is not None
+    assert spec.name == "main"
+
+
+@pytest.mark.usefixtures("import_hook")
+def test_import():
+    from .test_module import main  # type: ignore[reportAttributeAccessIssue]
+
+    assert str(main.hello()) == """\
+<h1>
+    Hello, World!
+</h1>"""

tests/test_module/main.px

@@ -0,0 +1,4 @@
+from pyjsx import jsx
+
+def hello():
+    return <h1>Hello, World!</h1>