Add --stub-extension CLI option (#142)

The option could be helpful for downstream doc-generation tools that
expect `.py` files only.

Author: sizmailov

README.md

@@ -28,8 +28,10 @@ pybind11-stubgen [-h]
                  [--ignore-invalid-identifiers REGEX]
                  [--ignore-unresolved-names REGEX]
                  [--ignore-all-errors]
-                 [--numpy-array-wrap-with-annotated-fixed-size| --numpy-array-remove-parameters]
+                 [--numpy-array-wrap-with-annotated|
+                  --numpy-array-remove-parameters]
                  [--print-invalid-expressions-as-is]
                  [--exit-code]
+                 [--stub-extension EXT]
                  MODULE_NAME
 ```

pybind11_stubgen/__init__.py

@@ -147,6 +147,16 @@ def regex(pattern_str: str) -> re.Pattern:
         help="Don't write stubs. Parses module and report errors",
     )
 
+    parser.add_argument(
+        "--stub-extension",
+        type=str,
+        default="pyi",
+        metavar="EXT",
+        choices=["pyi", "py"],
+        help="The file extension of the generated stubs. "
+        "Must be 'pyi' (default) or 'py'",
+    )
+
     parser.add_argument(
         "module_name",
         metavar="MODULE_NAME",
@@ -246,6 +256,7 @@ def main():
         out_dir,
         sub_dir=sub_dir,
         dry_run=args.dry_run,
+        writer=Writer(stub_ext=args.stub_extension),
     )
 
 
@@ -274,6 +285,7 @@ def run(
     out_dir: Path,
     sub_dir: Path | None,
     dry_run: bool,
+    writer: Writer,
 ):
     module = parser.handle_module(
         QualifiedName.from_str(module_name), importlib.import_module(module_name)
@@ -286,8 +298,6 @@ def run(
     if dry_run:
         return
 
-    writer = Writer()
-
     out_dir.mkdir(exist_ok=True, parents=True)
     writer.write_module(module, printer, to=out_dir, sub_dir=sub_dir)
 

pybind11_stubgen/writer.py

@@ -7,6 +7,10 @@
 
 
 class Writer:
+    def __init__(self, stub_ext: str = "pyi"):
+        # assert stub_extension in ["pyi", "py"]
+        self.stub_ext: str = stub_ext
+
     def write_module(
         self, module: Module, printer: Printer, to: Path, sub_dir: Path | None = None
     ):
@@ -17,10 +21,9 @@ def write_module(
                 sub_dir = Path(module.name)
             module_dir = to / sub_dir
             module_dir.mkdir(exist_ok=True)
-            module_file = module_dir / "__init__.pyi"
+            module_file = module_dir / f"__init__.{self.stub_ext}"
         else:
-            module_file = to / f"{module.name}.pyi"
-
+            module_file = to / f"{module.name}.{self.stub_ext}"
         with open(module_file, "w", encoding="utf-8") as f:
             f.writelines(line + "\n" for line in printer.print_module(module))