klayout/scripts/stubgen.py

477 lines
15 KiB
Python

""" Stub file generation routines.
This module contains routines to generate stub files from klayout's python API.
This uses the `tl` module of the API, which offers an introspection layer to
the C-extension modules.
"""
from collections import Counter
from copy import copy
from dataclasses import dataclass, field
from functools import wraps
import functools
from sys import argv
from textwrap import indent
from typing import Any, List, Optional, Tuple, Union
import pya # initialize all modules
import klayout.tl as ktl
def qualified_name(_class: ktl.Class) -> str:
name = _class.name()
if _class.parent():
return f"{qualified_name(_class.parent())}.{name}"
else:
return name
def superclass(_class: ktl.Class) -> str:
if _class.base():
return superclass(_class.base())
else:
return _class.name()
def is_reserved_word(name: str) -> bool:
wordlist = [
"and",
"del",
"from",
"not",
"while",
"as",
"elif",
"global",
"or",
"with",
"assert",
"else",
"if",
"pass",
"yield",
"break",
"except",
"import",
"print",
"class",
"exec",
"in",
"raise",
"continue",
"finally",
"is",
"return",
"def",
"for",
"lambda",
"try",
"None",
]
return name in wordlist
_type_dict = dict()
_type_dict[ktl.ArgType.TypeBool] = "bool"
_type_dict[ktl.ArgType.TypeChar] = "str"
_type_dict[ktl.ArgType.TypeDouble] = "float"
_type_dict[ktl.ArgType.TypeFloat] = "float"
_type_dict[ktl.ArgType.TypeInt] = "int"
_type_dict[ktl.ArgType.TypeLong] = "int"
_type_dict[ktl.ArgType.TypeLongLong] = "int"
# _type_dict[ktl.ArgType.TypeMap] = None
# _type_dict[ktl.ArgType.TypeObject] = None
_type_dict[ktl.ArgType.TypeSChar] = "str"
_type_dict[ktl.ArgType.TypeShort] = "int"
_type_dict[ktl.ArgType.TypeString] = "str"
_type_dict[ktl.ArgType.TypeByteArray] = "bytes"
_type_dict[ktl.ArgType.TypeUChar] = "str"
_type_dict[ktl.ArgType.TypeUInt] = "int"
_type_dict[ktl.ArgType.TypeULong] = "int"
_type_dict[ktl.ArgType.TypeULongLong] = "int"
_type_dict[ktl.ArgType.TypeUShort] = "int"
_type_dict[ktl.ArgType.TypeVar] = "Any"
# _type_dict[ktl.ArgType.TypeVector] = None
_type_dict[ktl.ArgType.TypeVoid] = "None"
_type_dict[ktl.ArgType.TypeVoidPtr] = "None"
def _translate_type(
arg_type: ktl.ArgType, within_class: ktl.Class, is_return=False
) -> str:
"""Translates klayout's C-type to a type in Python.
This function is equivalent to the `type_to_s` in `pyaModule.cc`.
See also `type_to_s` in `layGSIHelpProvider.cc`"""
py_str: str = ""
if arg_type.type() == ktl.ArgType.TypeObject:
if within_class.module() and arg_type.cls().module() != within_class.module():
py_str = arg_type.cls().module() + "." + qualified_name(arg_type.cls())
else:
py_str = qualified_name(arg_type.cls())
elif arg_type.type() == ktl.ArgType.TypeMap:
inner_key = _translate_type(arg_type.inner_k(), within_class, is_return)
inner_val = _translate_type(arg_type.inner(), within_class, is_return)
py_str = f"Dict[{inner_key}, {inner_val}]"
elif arg_type.type() == ktl.ArgType.TypeVector:
if is_return:
py_str = f"List[{_translate_type(arg_type.inner(), within_class, is_return)}]"
else:
py_str = f"Sequence[{_translate_type(arg_type.inner(), within_class, is_return)}]"
else:
py_str = _type_dict[arg_type.type()]
if arg_type.is_iter():
py_str = f"Iterator[{py_str}]"
if arg_type.has_default():
py_str = f"Optional[{py_str}] = ..."
return py_str
@dataclass
class Stub:
signature: str
name: Any
docstring: str
indent_docstring: bool = True
child_stubs: List["Stub"] = field(default_factory=list)
decorator: str = ""
def __eq__(self, __o: object) -> bool:
if not isinstance(__o, Stub):
return False
return (
self.signature == __o.signature
and self.child_stubs == __o.child_stubs
and self.decorator == __o.decorator
)
def __lt__(self, other: "Stub") -> bool:
# mypy complains if an overload with float happens before int
self_signature = self.signature.replace(": int", "0").replace(": float", "1")
other_signature = other.signature.replace(": int", "0").replace(": float", "1")
self_sortkey = self.name, len(self.signature.split(",")), self_signature
other_sortkey = other.name, len(other.signature.split(",")), other_signature
return self_sortkey < other_sortkey
def __hash__(self):
return hash(self.format_stub(include_docstring=False))
def format_stub(self, include_docstring=True):
lines = []
lines.extend(self.decorator.splitlines())
if self.indent_docstring: # all but properties
if include_docstring or len(self.child_stubs) > 0:
lines.append(self.signature + ":")
else:
lines.append(self.signature + ": ...")
else:
lines.append(self.signature)
stub_str = "\n".join(lines)
lines = []
lines.append('r"""')
lines.extend(self.docstring.splitlines())
lines.append('"""')
doc_str = "\n".join(lines)
# indent only if it is required (methods, not properties)
if self.indent_docstring:
doc_str = indent(doc_str, " " * 4)
if include_docstring:
stub_str += "\n"
stub_str += doc_str
for stub in self.child_stubs:
stub_str += "\n"
stub_str += indent(
stub.format_stub(include_docstring=include_docstring), " " * 4
)
if self.indent_docstring and (include_docstring or len(self.child_stubs)):
stub_str += "\n"
stub_str += indent("...", " " * 4)
return stub_str
@dataclass(eq=False)
class MethodStub(Stub):
indent_docstring: bool = True
@dataclass(eq=False)
class PropertyStub(Stub):
indent_docstring: bool = False
@dataclass(eq=False)
class ClassStub(Stub):
indent_docstring: bool = True
def get_child_classes(c: ktl.Class):
child_classes = []
for c_child in c.each_child_class():
child_classes.append(c_child)
return sorted(child_classes, key=lambda cls: cls.name())
def get_py_child_classes(c: ktl.Class):
for c_child in get_child_classes(c):
yield c_child
def get_py_methods(
c: ktl.Class,
) -> List[Stub]:
translate_arg_type = functools.partial(_translate_type, within_class=c, is_return=False)
translate_ret_type = functools.partial(_translate_type, within_class=c, is_return=True)
def _get_arglist(
m: ktl.Method, self_str: str
) -> List[Tuple[str, Optional[ktl.ArgType]]]:
args: List[Tuple[str, Optional[ktl.ArgType]]] = [(self_str, None)]
for i, a in enumerate(m.each_argument()):
argname = a.name()
if is_reserved_word(argname):
argname += "_"
elif not argname:
argname = f"arg{i}"
args.append((argname, a))
return args
def _format_args(arglist: List[Tuple[str, Optional[str]]]):
args = []
for argname, argtype in arglist:
if argtype:
args.append(f"{argname}: {argtype}")
else:
args.append(argname)
return ", ".join(args)
def format_args(m: ktl.Method, self_str: str = "self") -> str:
arg_list = _get_arglist(m, self_str=self_str)
new_arglist: List[Tuple[str, Optional[str]]] = []
for argname, a in arg_list:
if a:
new_arglist.append((argname, translate_arg_type(a)))
else:
new_arglist.append((argname, None))
return _format_args(new_arglist)
# Collect all properties here
properties: List[Stub] = list()
# Extract all instance properties
for f in c.python_properties(False):
name = f.getter().name()
getter = None
if len(f.getter().methods()) > 0:
getter = f.getter().methods()[0]
setter = None
if len(f.setter().methods()) > 0:
setter = f.setter().methods()[0]
if getter and setter:
# Full property
ret_type = translate_ret_type(getter.ret_type())
doc = "Getter:\n" + getter.doc() + "\nSetter:\n" + setter.doc()
properties.append(
PropertyStub(
decorator="",
signature=f"{name}: {ret_type}",
name=name,
docstring=doc,
)
)
elif getter:
# Only getter
ret_type = translate_ret_type(getter.ret_type())
doc = getter.doc()
properties.append(
MethodStub(
decorator="@property",
signature=f"def {name}(self) -> {ret_type}",
name=name,
docstring=doc,
)
)
elif setter:
# Only setter
doc = "WARNING: This variable can only be set, not retrieved.\n" + setter.doc()
properties.append(
MethodStub(
decorator="@property",
signature=f"def {name}(self) -> None",
name=name,
docstring=doc,
)
)
# Extract all class properties (TODO: setters not supported currently)
for f in c.python_properties(True):
name = f.getter().name()
if len(f.getter().methods()) > 0:
getter = f.getter().methods()[0]
ret_type = translate_ret_type(getter.ret_type())
doc = getter.doc()
properties.append(
PropertyStub(
decorator="",
signature=f"{name}: ClassVar[{ret_type}]",
name=name,
docstring=doc,
)
)
# Extract all classmethods
classmethods: List[Stub] = list()
for f in c.python_methods(True):
name = f.name()
decorator = ""
if len(f.methods()) > 1:
decorator = "@overload\n"
decorator += "@classmethod"
for m in f.methods():
ret_type = translate_ret_type(m.ret_type())
classmethods.append(
MethodStub(
decorator=decorator,
signature=f"def {name}({format_args(m, 'cls')}) -> {ret_type}",
name=name,
docstring=m.doc(),
)
)
# Extract bound methods
boundmethods: List[Stub] = list()
for f in c.python_methods(False):
name = f.name()
decorator = ""
if len(f.methods()) > 1:
decorator = "@overload\n"
for m in f.methods():
if name == "__init__":
ret_type = "None"
else:
ret_type = translate_ret_type(m.ret_type())
arg_list = _get_arglist(m, "self")
# TODO: fix type errors
# Exceptions:
# For X.__eq__(self, a:X), treat second argument as type object instead of X
if name in ("__eq__", "__ne__"):
if arg_list[1][1].cls() is not None and arg_list[1][1].cls().__gsi_id__ == c.__gsi_id__:
arg_list[1] = arg_list[1][0], "object"
# X._assign(self, other:X), mypy complains if _assign is defined at base class.
# We can't specialize other in this case.
elif name in ("_assign", "assign"):
assert arg_list[1][1] is not None
assert arg_list[1][1].type() == ktl.ArgType.TypeObject
arg_list[1] = arg_list[1][0], superclass(arg_list[1][1].cls())
else:
new_arg_list = []
for argname, a in arg_list:
if a:
new_arg_list.append((argname, translate_arg_type(a)))
else:
new_arg_list.append((argname, a))
arg_list = new_arg_list
formatted_args = _format_args(arg_list)
boundmethods.append(
MethodStub(
decorator=decorator,
signature=f"def {name}({formatted_args}) -> {ret_type}",
name=name,
docstring=m.doc(),
)
)
boundmethods = sorted(boundmethods, key=lambda m: m.signature)
properties = sorted(properties, key=lambda m: m.signature)
classmethods = sorted(classmethods, key=lambda m: m.signature)
return_list: List[Stub] = properties + classmethods + boundmethods
return return_list
def get_class_stub(
c: ktl.Class,
ignore: List[ktl.Class] = None,
module: str = "",
) -> ClassStub:
base = ""
if c.base():
base = f"({c.base().name()})"
if c.module() != module:
full_name = c.module() + "." + c.name()
else:
full_name = c.name()
_cstub = ClassStub(
signature="class " + full_name + base, docstring=c.doc(), name=full_name
)
child_attributes = get_py_methods(c)
for child_c in get_child_classes(c):
_cstub.child_stubs.append(
get_class_stub(
child_c,
ignore=ignore,
module=c.module(),
)
)
for stub in child_attributes:
_cstub.child_stubs.append(stub)
return _cstub
def get_classes(module: str) -> List[ktl.Class]:
_classes = []
for c in ktl.Class.each_class():
if c.module() != module:
continue
_classes.append(c)
return sorted(_classes, key=lambda cls: cls.name())
def get_module_stubs(module: str) -> List[ClassStub]:
_stubs = []
_classes = get_classes(module)
for c in _classes:
_cstub = get_class_stub(c, ignore=_classes, module=module)
_stubs.append(_cstub)
return _stubs
def print_mod(module, dependencies):
print("from typing import Any, ClassVar, Dict, Sequence, List, Iterator, Optional")
print("from typing import overload")
print("from __future__ import annotations")
for dep in dependencies:
print(f"import klayout.{dep} as {dep}")
for stub in get_module_stubs(module):
print(stub.format_stub(include_docstring=True) + "\n")
if __name__ == "__main__":
if len(argv) < 2:
print("Specify module in argument")
exit(1)
if len(argv) == 2:
print_mod(argv[1], [])
else:
print_mod(argv[1], argv[2].split(","))