def emit_node(
self,
name: str,
type_: NodeType,
properties: Optional[NodeProperties] = None,
) -> None:
"""Create a new node. Nodes can be classes, packages, participants etc."""
if properties is None:
properties = NodeProperties(label=name)
stereotype = " << interface >>" if type_ is NodeType.INTERFACE else ""
nodetype = self.NODES[type_]
if properties.color and properties.color != self.DEFAULT_COLOR:
color = f" #{properties.color}"
else:
color = ""
body = []
if properties.attrs:
body.extend(properties.attrs)
if properties.methods:
for func in properties.methods:
args = self._get_method_arguments(func)
line = f"{func.name}({', '.join(args)})"
if func.returns:
line += " -> " + get_annotation_label(func.returns)
body.append(line)
label = properties.label if properties.label is not None else name
if properties.fontcolor and properties.fontcolor != self.DEFAULT_COLOR:
label = f"<color:{properties.fontcolor}>{label}</color>"
self.emit(f'{nodetype} "{label}" as {name}{stereotype}{color} {{')
self._inc_indent()
for line in body:
self.emit(line)
self._dec_indent()
self.emit("}")
def emit_node(
self,
name: str,
type_: NodeType,
properties: Optional[NodeProperties] = None,
) -> None:
"""Create a new node.
Nodes can be classes, packages, participants etc.
"""
if properties is None:
properties = NodeProperties(label=name)
stereotype = "~~Interface~~" if type_ is NodeType.INTERFACE else ""
nodetype = self.NODES[type_]
body = []
if properties.attrs:
body.extend(properties.attrs)
if properties.methods:
for func in properties.methods:
args = self._get_method_arguments(func)
line = f"{func.name}({', '.join(args)})"
if func.returns:
line += " -> " + get_annotation_label(func.returns)
body.append(line)
name = name.split(".")[-1]
self.emit(f"{nodetype} {name}{stereotype} {{")
self._inc_indent()
for line in body:
self.emit(line)
self._dec_indent()
self.emit("}")
def _build_label_for_node(self,
properties: NodeProperties,
is_interface: Optional[bool] = False) -> str:
if not properties.label:
return ""
label: str = properties.label
if is_interface:
# add a stereotype
label = "<<interface>>\\n" + label
if properties.attrs is None and properties.methods is None:
# return a "compact" form which only displays the class name in a box
return label
# Add class attributes
attrs: List[str] = properties.attrs or []
label = "{" + label + "|" + r"\l".join(attrs) + r"\l|"
# Add class methods
methods: List[nodes.FunctionDef] = properties.methods or []
for func in methods:
args = self._get_method_arguments(func)
label += fr"{func.name}({', '.join(args)})"
if func.returns:
label += ": " + get_annotation_label(func.returns)
label += r"\l"
label += "}"
return label
def _get_method_arguments(method: nodes.FunctionDef) -> List[str]:
if method.args.args is None:
return []
first_arg = 0 if method.type in {"function", "staticmethod"} else 1
arguments: List[nodes.AssignName] = method.args.args[first_arg:]
annotations = dict(zip(arguments, method.args.annotations[first_arg:]))
for arg in arguments:
annotation_label = ""
ann = annotations.get(arg)
if ann:
annotation_label = get_annotation_label(ann)
annotations[arg] = annotation_label
return [
f"{arg.name}: {ann}" if ann else f"{arg.name}"
for arg, ann in annotations.items()
]
def _get_method_arguments(method: nodes.FunctionDef) -> List[str]:
if method.args.args:
arguments: List[nodes.AssignName] = [
arg for arg in method.args.args if arg.name != "self"
]
else:
arguments = []
annotations = dict(zip(arguments, method.args.annotations[1:]))
for arg in arguments:
annotation_label = ""
ann = annotations.get(arg)
if ann:
annotation_label = get_annotation_label(ann)
annotations[arg] = annotation_label
return [
f"{arg.name}: {ann}" if ann else f"{arg.name}"
for arg, ann in annotations.items()
]
def get_values(self, obj):
"""get label and shape for classes.
The label contains all attributes and methods
"""
label = obj.title
if obj.shape == "interface":
label = "«interface»\\n%s" % label
if not self.config.only_classnames:
label = r"{}|{}\l|".format(label, r"\l".join(obj.attrs))
for func in obj.methods:
return_type = (
f": {get_annotation_label(func.returns)}" if func.returns else ""
)
if func.args.args:
args = [arg for arg in func.args.args if arg.name != "self"]
else:
args = []
annotations = dict(zip(args, func.args.annotations[1:]))
for arg in args:
annotation_label = ""
ann = annotations.get(arg)
if ann:
annotation_label = get_annotation_label(ann)
annotations[arg] = annotation_label
args = ", ".join(
f"{arg.name}: {ann}" if ann else f"{arg.name}"
for arg, ann in annotations.items()
)
label = fr"{label}{func.name}({args}){return_type}\l"
label = "{%s}" % label
if is_exception(obj.node):
return dict(fontcolor="red", label=label, shape="record")
return dict(label=label, shape="record")
def _build_label_for_node(self, properties: NodeProperties) -> str:
if not properties.label:
return ""
label: str = properties.label
if properties.attrs is None and properties.methods is None:
# return a "compact" form which only displays the class name in a box
return label
# Add class attributes
attrs: list[str] = properties.attrs or []
attrs_string = r"\l".join(attr.replace("|", r"\|") for attr in attrs)
label = rf"{{{label}|{attrs_string}\l|"
# Add class methods
methods: list[nodes.FunctionDef] = properties.methods or []
for func in methods:
args = self._get_method_arguments(func)
label += rf"{func.name}({', '.join(args)})"
if func.returns:
label += ": " + get_annotation_label(func.returns)
label += r"\l"
label += "}"
return label
def test_get_annotation_label_of_return_type(node_text: str,
expected_label: str) -> None:
func = astroid.extract_node(node_text)
assert isinstance(func, nodes.FunctionDef)
assert get_annotation_label(func.returns) == expected_label
