Support converting unknown Python scripts to marimo notebooks

Enable `marimo convert script.py` to convert regular Python scripts.

Author: manzt

marimo/_cli/convert/commands.py

@@ -28,9 +28,9 @@ def convert(
     filename: str,
     output: Optional[Path],
 ) -> None:
-    r"""Convert a Jupyter notebook or Markdown file to a marimo notebook.
+    r"""Convert a Jupyter notebook, Markdown file, or unknown Python script to a marimo notebook.
 
-    The argument may be either a path to a local .ipynb/.md file,
+    The argument may be either a path to a local .ipynb/.md/.py file,
     or an .ipynb/.md file hosted on GitHub.
 
     Example usage:
@@ -41,8 +41,18 @@ def convert(
 
         marimo convert your_nb.md -o your_nb.py
 
+    or
+
+        marimo convert script.py -o your_nb.py
+
     Jupyter notebook conversion will strip out all outputs. Markdown cell
-    conversion with occur on the presence of `{python}` code fences.
+    conversion will occur on the presence of `{python}` code fences.
+
+    For .py files:
+    - If the file is already a valid marimo notebook, no conversion is performed
+    - Unknown Python scripts are converted by preserving the header (docstrings/comments)
+      and splitting the code into cells, with __main__ blocks separated
+
     After conversion, you can open the notebook in the editor:
 
         marimo edit your_nb.py
@@ -53,15 +63,41 @@ def convert(
     """
 
     ext = Path(filename).suffix
-    if ext not in (".ipynb", ".md", ".qmd"):
-        raise click.UsageError("File must be an .ipynb or .md file")
+    if ext not in (".ipynb", ".md", ".qmd", ".py"):
+        raise click.UsageError("File must be an .ipynb, .md, or .py file")
 
     text = load_external_file(filename, ext)
     if ext == ".ipynb":
         notebook = MarimoConvert.from_ipynb(text).to_py()
-    else:
-        assert ext in (".md", ".qmd")
+    elif ext in (".md", ".qmd"):
         notebook = MarimoConvert.from_md(text).to_py()
+    else:
+        assert ext == ".py"
+        # First check if it's already a valid marimo notebook
+        from marimo._ast.parse import parse_notebook
+
+        try:
+            parsed = parse_notebook(text)
+        except SyntaxError:
+            # File has syntax errors
+            echo("File cannot be converted. It may have syntax errors.")
+            return
+
+        if parsed and parsed.valid:
+            # Already a valid marimo notebook
+            echo("File is already a valid marimo notebook.")
+            return
+
+        # Check if it has the violation indicating it's an unknown Python script
+        if parsed and any(
+            v.description == "Unknown content beyond header"
+            for v in parsed.violations
+        ):
+            notebook = MarimoConvert.from_unknown_py_script(text).to_py()
+        else:
+            # File has other issues (syntax errors, etc.)
+            echo("File cannot be converted. It may have syntax errors.")
+            return
 
     if output:
         output_path = Path(output)

marimo/_convert/converters.py

@@ -49,6 +49,22 @@ def from_py(source: str) -> MarimoConverterIntermediate:
         ir = parse_notebook(source) or EMPTY_NOTEBOOK_SERIALIZATION
         return MarimoConverterIntermediate(ir)
 
+    @staticmethod
+    def from_unknown_py_script(source: str) -> MarimoConverterIntermediate:
+        """Convert from unknown Python script to marimo notebook.
+
+        This should only be used when the .py file is not already a valid
+        marimo notebook.
+
+        Args:
+            source: Unknown Python script source code string
+        """
+        from marimo._convert.unknown_python import (
+            convert_unknown_py_to_notebook_ir,
+        )
+
+        return MarimoConvert.from_ir(convert_unknown_py_to_notebook_ir(source))
+
     @staticmethod
     def from_md(source: str) -> MarimoConverterIntermediate:
         """Convert from markdown source code.

marimo/_convert/unknown_python.py

@@ -0,0 +1,91 @@
+# Copyright 2024 Marimo. All rights reserved.
+"""Convert unknown Python scripts to marimo notebooks."""
+
+from __future__ import annotations
+
+from marimo._ast.parse import Parser
+from marimo._schemas.serialization import (
+    AppInstantiation,
+    CellDef,
+    NotebookSerialization,
+)
+
+
+def convert_unknown_py_to_notebook_ir(source: str) -> NotebookSerialization:
+    """Convert an unknown Python script to marimo notebook IR.
+
+    This should only be called after verifying the file is not already
+    a valid marimo notebook. It converts by:
+    1. Preserving the header (docstrings/comments)
+    2. Splitting on 'if __name__ == "__main__":' if present:
+       - Everything before goes into the first cell
+       - The main block is transformed into a second cell with:
+         - 'if __name__ == "__main__":' replaced with 'def _main_():'
+         - '_main_()' called at the end
+    3. If no main block exists, all content goes into a single cell
+    """
+    parser = Parser(source)
+    body = parser.node_stack()
+    cells: list[CellDef] = []
+    header_result = parser.parse_header(body)
+
+    header = header_result.unwrap()
+    remaining = parser.extractor.contents[len(header.value) :].strip()
+
+    if remaining:
+        # Split on if __name__ == "__main__":
+        main_pattern = 'if __name__ == "__main__":'
+        if main_pattern in remaining:
+            parts = remaining.split(main_pattern, 1)
+            before_main = parts[0].strip()
+
+            # Replace the if __name__ == "__main__": with def _main_():
+            main_block = "def _main_():" + parts[1] + "\n\n_main_()"
+
+            # Create cells
+            if before_main:
+                cells.append(
+                    CellDef(
+                        lineno=1,
+                        col_offset=0,
+                        end_lineno=1,
+                        end_col_offset=0,
+                        code=before_main,
+                        name="_",
+                        options={},
+                    )
+                )
+
+            cells.append(
+                CellDef(
+                    lineno=2,
+                    col_offset=0,
+                    end_lineno=2,
+                    end_col_offset=0,
+                    code=main_block,
+                    name="_",
+                    options={},
+                )
+            )
+        else:
+            # No main block, put everything in one cell
+            cells.append(
+                CellDef(
+                    lineno=1,
+                    col_offset=0,
+                    end_lineno=1,
+                    end_col_offset=0,
+                    code=remaining,
+                    name="_",
+                    options={},
+                )
+            )
+
+    return NotebookSerialization(
+        header=header,
+        version=None,
+        app=AppInstantiation(),
+        cells=cells,
+        violations=header_result.violations,
+        valid=False,
+    )

tests/_cli/snapshots/python_script_no_main_to_marimo.txt

@@ -0,0 +1,20 @@
+"""Simple calculation script."""
+
+import marimo
+
+
+app = marimo.App()
+
+
+@app.cell
+def _():
+    x = 5
+    y = 10
+    result = x + y
+    print(f"Result: {result}")
+    return
+
+
+if __name__ == "__main__":
+    app.run()
+

tests/_cli/snapshots/python_script_to_marimo.txt

@@ -0,0 +1,30 @@
+"""A simple Python script."""
+
+import marimo
+
+
+app = marimo.App()
+
+
+@app.cell
+def _():
+    import sys
+
+    def main():
+        print("Hello, World!")
+    return main, sys
+
+
+@app.cell
+def _(main, sys):
+    def _main_():
+        main()
+        sys.exit(0)
+
+    _main_()
+    return
+
+
+if __name__ == "__main__":
+    app.run()
+

tests/_cli/test_cli_convert.py

@@ -167,7 +167,7 @@ def test_convert_invalid_file(tmp_path: Path) -> None:
             text=True,
         )
         assert p.returncode != 0
-        assert "File must be an .ipynb or .md file" in p.stderr
+        assert "File must be an .ipynb, .md, or .py file" in p.stderr
 
     @staticmethod
     def test_convert_remote_ipynb(http_server: MockHTTPServer) -> None:
@@ -255,7 +255,7 @@ def test_convert_remote_invalid_file(http_server: MockHTTPServer) -> None:
             text=True,
         )
         assert p.returncode != 0
-        assert "File must be an .ipynb or .md file" in p.stderr
+        assert "File must be an .ipynb, .md, or .py file" in p.stderr
 
     @staticmethod
     def test_convert_nonexistent_remote_file(
@@ -274,3 +274,82 @@ def test_convert_nonexistent_remote_file(
         assert p.returncode != 0
         # The error message will be from urllib.error.HTTPError
         assert "HTTP Error 404" in p.stderr or "Not Found" in p.stderr
+
+    @staticmethod
+    def test_convert_existing_marimo_notebook(tmp_path: Path) -> None:
+        """Test that converting an existing marimo notebook prints a message."""
+        marimo_path = tmp_path / "existing_marimo.py"
+        marimo_content = """import marimo
+
+__generated_with = "0.10.0"
+app = marimo.App()
+
+
+@app.cell
+def __():
+    print("Hello from marimo!")
+    return
+
+
+if __name__ == "__main__":
+    app.run()
+"""
+        marimo_path.write_text(marimo_content)
+
+        p = subprocess.run(
+            ["marimo", "convert", str(marimo_path)],
+            capture_output=True,
+            text=True,
+        )
+        assert p.returncode == 0
+        assert "File is already a valid marimo notebook." in p.stdout
+
+    @staticmethod
+    def test_convert_unknown_python_script(tmp_path: Path) -> None:
+        """Test converting an unknown Python script."""
+        script_path = tmp_path / "script.py"
+        script_content = '''"""A simple Python script."""
+
+import sys
+
+def main():
+    print("Hello, World!")
+
+if __name__ == "__main__":
+    main()
+    sys.exit(0)
+'''
+        script_path.write_text(script_content)
+
+        p = subprocess.run(
+            ["marimo", "convert", str(script_path)],
+            capture_output=True,
+            text=True,
+        )
+        assert p.returncode == 0, p.stderr
+        output = p.stdout
+        output = re.sub(r"__generated_with = .*", "", output)
+        snapshot("python_script_to_marimo.txt", output)
+
+    @staticmethod
+    def test_convert_python_script_no_main(tmp_path: Path) -> None:
+        """Test converting a Python script without main block."""
+        script_path = tmp_path / "simple_script.py"
+        script_content = '''"""Simple calculation script."""
+
+x = 5
+y = 10
+result = x + y
+print(f"Result: {result}")
+'''
+        script_path.write_text(script_content)
+
+        p = subprocess.run(
+            ["marimo", "convert", str(script_path)],
+            capture_output=True,
+            text=True,
+        )
+        assert p.returncode == 0, p.stderr
+        output = p.stdout
+        output = re.sub(r"__generated_with = .*", "", output)
+        snapshot("python_script_no_main_to_marimo.txt", output)

tests/_convert/snapshots/empty_main_block.py.txt

@@ -0,0 +1,25 @@
+"""Script with empty main."""
+
+import marimo
+
+__generated_with = "0.0.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+    data = [1, 2, 3, 4, 5]
+    return
+
+
+@app.cell
+def _():
+    def _main_():
+        pass
+
+    _main_()
+    return
+
+
+if __name__ == "__main__":
+    app.run()

tests/_convert/snapshots/minimal_script.py.txt

@@ -0,0 +1,16 @@
+import marimo
+
+__generated_with = "0.0.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+    x = 5
+    y = 10
+    print(x + y)
+    return
+
+
+if __name__ == "__main__":
+    app.run()