in metaflow/cmd/develop/stub_generator.py [0:0]
def _generate_class_stub(self, name: str, clazz: type) -> str:
debug.stubgen_exec("Generating class stub for %s" % name)
skip_init = issubclass(clazz, (TriggeredRun, DeployedFlow))
if issubclass(clazz, DeployerImpl):
if clazz.TYPE is not None:
clazz_type = clazz.TYPE.replace("-", "_")
self._deployer_injected_methods.setdefault(clazz_type, {})[
"deployer"
] = (self._current_module_name + "." + name)
buff = StringIO()
# Class prototype
buff.write("class " + name.split(".")[-1] + "(")
# Add super classes
for c in clazz.__bases__:
name_with_module = self._get_element_name_with_module(c, force_import=True)
buff.write(name_with_module + ", ")
# Add metaclass
name_with_module = self._get_element_name_with_module(
clazz.__class__, force_import=True
)
buff.write("metaclass=" + name_with_module + "):\n")
# Add class docstring
if clazz.__doc__:
buff.write('%s"""\n' % TAB)
my_doc = inspect.cleandoc(clazz.__doc__)
init_blank = True
for line in my_doc.split("\n"):
if init_blank and len(line.strip()) == 0:
continue
init_blank = False
buff.write("%s%s\n" % (TAB, line.rstrip()))
buff.write('%s"""\n' % TAB)
# For NamedTuple, we have __annotations__ but no __init__. In that case,
# we are going to "create" a __init__ function with the annotations
# to show what the class takes.
annotation_dict = None
init_func = None
for key, element in clazz.__dict__.items():
func_deco = None
if isinstance(element, staticmethod):
func_deco = "@staticmethod"
element = element.__func__
elif isinstance(element, classmethod):
func_deco = "@classmethod"
element = element.__func__
if key == "__init__":
if skip_init:
continue
init_func = element
elif key == "__annotations__":
annotation_dict = element
if inspect.isfunction(element):
if not element.__name__.startswith("_") or element.__name__.startswith(
"__"
):
if (
clazz == Deployer
and element.__name__ in self._deployer_injected_methods
):
# This is a method that was injected. It has docs but we need
# to parse it to generate the proper signature
func_doc = inspect.cleandoc(element.__doc__)
docs = split_docs(
func_doc,
[
("func_doc", StartEnd(0, 0)),
(
"param_doc",
param_section_header.search(func_doc)
or StartEnd(len(func_doc), len(func_doc)),
),
(
"return_doc",
return_section_header.search(func_doc)
or StartEnd(len(func_doc), len(func_doc)),
),
],
)
parameters, _ = parse_params_from_doc(docs["param_doc"])
return_type = self._deployer_injected_methods[element.__name__][
"deployer"
]
buff.write(
self._generate_function_stub(
key,
element,
sign=[
inspect.Signature(
parameters=[
inspect.Parameter(
"self",
inspect.Parameter.POSITIONAL_OR_KEYWORD,
)
]
+ parameters,
return_annotation=return_type,
)
],
indentation=TAB,
deco=func_deco,
)
)
elif (
clazz == DeployedFlow and element.__name__ == "from_deployment"
):
# We simply update the signature to list the return
# type as a union of all possible deployers
func_doc = inspect.cleandoc(element.__doc__)
docs = split_docs(
func_doc,
[
("func_doc", StartEnd(0, 0)),
(
"param_doc",
param_section_header.search(func_doc)
or StartEnd(len(func_doc), len(func_doc)),
),
(
"return_doc",
return_section_header.search(func_doc)
or StartEnd(len(func_doc), len(func_doc)),
),
],
)
parameters, _ = parse_params_from_doc(docs["param_doc"])
def _create_multi_type(*l):
return typing.Union[l]
all_types = [
v["from_deployment"][0]
for v in self._deployer_injected_methods.values()
]
if len(all_types) > 1:
return_type = _create_multi_type(*all_types)
else:
return_type = all_types[0] if len(all_types) else None
buff.write(
self._generate_function_stub(
key,
element,
sign=[
inspect.Signature(
parameters=[
inspect.Parameter(
"cls",
inspect.Parameter.POSITIONAL_OR_KEYWORD,
)
]
+ parameters,
return_annotation=return_type,
)
],
indentation=TAB,
doc=docs["func_doc"]
+ "\n\nParameters\n----------\n"
+ docs["param_doc"]
+ "\n\nReturns\n-------\n"
+ "%s\nA `DeployedFlow` object" % str(return_type),
deco=func_deco,
)
)
elif (
clazz == DeployedFlow
and element.__name__.startswith("from_")
and element.__name__[5:] in self._deployer_injected_methods
):
# Get the doc from the from_deployment method stored in
# self._deployer_injected_methods
func_doc = inspect.cleandoc(
self._deployer_injected_methods[element.__name__[5:]][
"from_deployment"
][1]
or ""
)
docs = split_docs(
func_doc,
[
("func_doc", StartEnd(0, 0)),
(
"param_doc",
param_section_header.search(func_doc)
or StartEnd(len(func_doc), len(func_doc)),
),
(
"return_doc",
return_section_header.search(func_doc)
or StartEnd(len(func_doc), len(func_doc)),
),
],
)
parameters, _ = parse_params_from_doc(docs["param_doc"])
return_type = self._deployer_injected_methods[
element.__name__[5:]
]["from_deployment"][0]
buff.write(
self._generate_function_stub(
key,
element,
sign=[
inspect.Signature(
parameters=[
inspect.Parameter(
"cls",
inspect.Parameter.POSITIONAL_OR_KEYWORD,
)
]
+ parameters,
return_annotation=return_type,
)
],
indentation=TAB,
doc=docs["func_doc"]
+ "\n\nParameters\n----------\n"
+ docs["param_doc"]
+ "\n\nReturns\n-------\n"
+ docs["return_doc"],
deco=func_deco,
)
)
else:
if (
issubclass(clazz, DeployedFlow)
and clazz.TYPE is not None
and key == "from_deployment"
):
clazz_type = clazz.TYPE.replace("-", "_")
# Record docstring for this function
self._deployer_injected_methods.setdefault(clazz_type, {})[
"from_deployment"
] = (
self._current_module_name + "." + name,
element.__doc__,
)
buff.write(
self._generate_function_stub(
key,
element,
indentation=TAB,
deco=func_deco,
)
)
elif isinstance(element, property):
if element.fget:
buff.write(
self._generate_function_stub(
key, element.fget, indentation=TAB, deco="@property"
)
)
if element.fset:
buff.write(
self._generate_function_stub(
key, element.fset, indentation=TAB, deco="@%s.setter" % key
)
)
# Special handling of classes that have injected methods
if clazz == Current:
# Multiple decorators can add the same object (trigger and trigger_on_finish)
# as examples so we sort it out.
resulting_dict = (
dict()
) # type Dict[str, List[inspect.Signature, str, List[str]]]
for deco_name, addl_current in self._addl_current.items():
for name, (sign, doc) in addl_current.items():
r = resulting_dict.setdefault(name, [sign, doc, []])
r[2].append("@%s" % deco_name)
for name, (sign, doc, decos) in resulting_dict.items():
buff.write(
self._generate_function_stub(
name,
sign=[sign],
indentation=TAB,
doc="(only in the presence of the %s decorator%s)\n\n"
% (", or ".join(decos), "" if len(decos) == 1 else "s")
+ doc,
deco="@property",
)
)
if not skip_init and init_func is None and annotation_dict:
buff.write(
self._generate_function_stub(
"__init__",
func=None,
sign=[
inspect.Signature(
parameters=[
inspect.Parameter(
name="self",
kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
)
]
+ [
inspect.Parameter(
name=name,
kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
annotation=annotation,
)
for name, annotation in annotation_dict.items()
]
)
],
indentation=TAB,
)
)
buff.write("%s...\n" % TAB)
return buff.getvalue()