feat(script)!: add a script to build stubs

change some doc

README.md

@@ -8,42 +8,88 @@ Add typing support for your YACS config by generating stub file.
 
 ## Install
 
-<details>
+Install from PyPI:
 
 ```sh
 pip install yacs-stubgen
 ```
 
-or install from this repo:
+<details>
+
+<summary>Other methods</summary>
+
+Install from this repo directly:
 
 ```sh
-pip install git+github.com/JamzumSum/yacs-stubgen.git
+pip install git+https://github.com/JamzumSum/yacs-stubgen.git
 ```
 
+Or you can download from our GitHub release and install package manually.
+
 </details>
 
 ## Usage
 
-Add typing support for your [yacs][yacs] config by appending just two lines:
+### Auto-Generate
+
+Add typing support for your [yacs][yacs] config by appending just three lines:
 
 ```py
+from yacs.config import CfgNode as CN
+
 _C.MODEL.DEVICE = 'cuda'
 ...
 # your config items above
 
-from yacs_stubgen import build_pyi
 # this line can be moved to the import header
-build_pyi(_C, __file__, var_name='_C')
-# _C is the CfgNode object, "_C" should be its name correctly
+from yacs_stubgen import build_pyi
+# this alias ensure you can import `AutoConfig` and use something like `isinstance`
+AutoConfig = CN
+# _C is the CfgNode object, "_C" should be its varname correctly
+# AutoConfig is an alias of CfgNode, "AutoConfig" should be its varname correctly
+build_pyi(_C, __file__, cls_name='AutoConfig', var_name='_C')
 ```
 
 **After** any run/import of this file, a stub file (*.pyi) will be generated.
 Then you will get typing and auto-complete support **if your IDE supports stub files**.
 
+Each time you change your config, you have to run/import this file again to apply the changes.
+
+### Build Script
+
+We have provided a script as an entrypoint. Simply run `yacstub <file/dir>` and it
+will generate stub file if one module contains a `CfgNode` object in global scope.
+
+```sh
+> yacstub ./conf    # specify a directory
+INFO: Generated conf/default.pyi
+> yacstub ./conf/default.py # specify a file
+INFO: Generated conf/default.pyi
+```
+
+Similarly, each time you change the config, you have to re-run the script to apply the changes.
+
+## How it works
+
+<details>
+
+**Stub files take precedence** in the case of both `filename.py` and `filename.pyi` exists.
+Once you pass in the config node, we will iterate over it and generate a stub file then save
+it as `filename.pyi` (that's why a path is required). Now supporting IDE will detect the stub
+file and is able to type-check and intellisense your code.
+
+However, the stub file does nothing with actual code executing. If you import the generated
+class (default as "AutoClass"), an `ImportError` will be raised. This time you can add an variable
+(aka. type alias) refers to `CfgNode` in the config file. We will override the type of this alias
+to our generated class ("AutoClass"). Thus you can import the "AutoClass" normally and intuitively,
+while the type alias is treated as "AutoClass" by IDE but is actually a `CfgNode` type.
+
+</details>
+
 ## License
 
 - [MIT](LICENSE)
-- [YACS][yacs] is under [Apache-2.0](https://github.com/rbgirshick/yacs/LICENSE)
+- [yacs][yacs] is under [Apache-2.0](https://github.com/rbgirshick/yacs/LICENSE)
 
 [yacs]: https://github.com/rbgirshick/yacs
 [home]: https://github.com/JamzumSum/yacs-stubgen

conf/default.py

@@ -16,4 +16,5 @@
 
 from yacs_stubgen import build_pyi
 
+AutoConfig = CN
 build_pyi(_C, __file__, var_name="_C")

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "yacs-stubgen"
-version = "0.1.2.post1"
+version = "0.2.0"
 description = "Generate stub file for yacs config."
 authors = ["JamzumSum <zzzzss990315@gmail.com>"]
 license = "MIT"
@@ -22,6 +22,9 @@ optional = true
 isort = "^5.10.1"
 black = "^22.8.0"
 
+[tool.poetry.scripts]
+yacstub = "genstub:main"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

src/genstub.py

@@ -0,0 +1,85 @@
+import argparse
+import importlib
+import logging
+import sys
+from pathlib import Path
+from typing import Iterable
+
+from yacs.config import CfgNode
+
+from yacs_stubgen import build_pyi
+
+logging.basicConfig(level="INFO", format="%(levelname)s: %(message)s")
+log = logging.getLogger(__name__)
+
+
+def inspect_module(mod_path: Path, ROOT: Path):
+    assert mod_path.exists()
+    rel_path = mod_path.relative_to(ROOT)
+    if rel_path.stem == "__init__":
+        rel_path = rel_path.parent
+    package = rel_path.with_suffix("").as_posix().replace("/", ".").replace("-", "_")
+
+    try:
+        mod = importlib.import_module(package)
+    except ImportError as e:
+        log.error(e)
+        return False
+
+    clsname = varname = ""
+    cfg = None
+    for k, v in mod.__dict__.items():
+        if isinstance(v, CfgNode):
+            cfg = v
+            varname = k
+        elif isinstance(v, type):
+            if issubclass(v, CfgNode):
+                clsname = k
+        if clsname and varname:
+            break
+
+    if cfg and varname:
+        if clsname:
+            build_pyi(cfg, mod_path, cls_name=clsname, var_name=varname)
+        else:
+            build_pyi(cfg, mod_path, var_name=varname)
+            log.warning(
+                "You haven't set an alias for CfgNode, the config class cannot be imported."
+            )
+        return True
+    return False
+
+
+def ibuild_pyi(pathes: Iterable[Path], ROOT: Path):
+    sys.path.insert(0, str(ROOT.absolute()))
+
+    try:
+        for py in pathes:
+            if any(i.startswith(".") for i in py.parts):
+                continue
+            if inspect_module(py, ROOT):
+                pyi = py.with_suffix(".pyi")
+                assert pyi.exists()
+                log.info(f"Generated {pyi.as_posix()}")
+    finally:
+        sys.path.pop(0)
+
+
+def rbuild_pyi(dir: Path):
+    return ibuild_pyi(dir.rglob("*.py"), dir)
+
+
+def main():
+    psr = argparse.ArgumentParser()
+    psr.add_argument("dir", nargs="?", default=".", type=Path)
+    args = psr.parse_args()
+
+    ROOT: Path = args.dir
+    if ROOT.is_dir():
+        rbuild_pyi(ROOT)
+    elif ROOT.is_file():
+        ibuild_pyi((ROOT,), ROOT.parent)
+
+
+if __name__ == "__main__":
+    main()

src/yacs_stubgen/__init__.py

@@ -5,13 +5,13 @@
 from yacs.config import CfgNode
 
 
-def to_py_obj(cfg: CfgNode, cls_name: str):
+def _to_py_obj(cfg: CfgNode, cls_name: str):
     classes = {}
     d = {}
     for k, v in cfg.items():
         if isinstance(v, CfgNode):
             _clsname = str.capitalize(k)
-            clss, _ = to_py_obj(v, _clsname)
+            clss, _ = _to_py_obj(v, _clsname)
             classes.update(clss)
             d[k] = _clsname
         else:
@@ -24,7 +24,15 @@ def to_py_obj(cfg: CfgNode, cls_name: str):
 def build_pyi(
     cfg: CfgNode, path: Union[Path, str], cls_name="AutoConfig", var_name="cfg"
 ):
-    d, _ = to_py_obj(cfg, cls_name)
+    """Generate a stub file (*.pyi) for the given config object.
+
+    :param cfg: the `CfgNode` object to generate stub.
+    :param path: the stub file output path or the cfg module file path. The suffix will be override.
+    :param cls_name: Generated name of the root config class. You can assign an alias of `CfgNode` to this param.
+    :param var_name: name of the `cfg` object. You should passin this param correctly.
+    """
+    assert cls_name != var_name, "class name should not be the same with var name"
+    d, _ = _to_py_obj(cfg, cls_name)
     d[var_name] = cls_name
     path = Path(path)
     with open(path.with_suffix(".pyi"), "w") as f: