Skip to content

Commit

Permalink
refactor: Change _elkjs typed dicts to pydantic elements
Browse files Browse the repository at this point in the history
  • Loading branch information
huyenngn committed Jun 5, 2024
1 parent 7fe3a45 commit 61cdead
Show file tree
Hide file tree
Showing 12 changed files with 338 additions and 333 deletions.
75 changes: 42 additions & 33 deletions capellambse_context_diagrams/_elkjs.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
from pathlib import Path

import capellambse
import typing_extensions as te
import pydantic

__all__ = [
"call_elkjs",
Expand Down Expand Up @@ -89,87 +89,96 @@
"""Options for labels to configure ELK layouting."""


class ELKInputData(te.TypedDict, total=False):
class BaseELKModel(pydantic.BaseModel):
"""Base class for ELK models."""

model_config = pydantic.ConfigDict(
extra="allow",
arbitrary_types_allowed=True,
)


class ELKInputData(BaseELKModel):
"""Data that can be fed to ELK."""

id: te.Required[str]
layoutOptions: LayoutOptions
children: cabc.MutableSequence[ELKInputChild] # type: ignore
edges: cabc.MutableSequence[ELKInputEdge]
id: str
layoutOptions: LayoutOptions = {}
children: cabc.MutableSequence[ELKInputChild] = []
edges: cabc.MutableSequence[ELKInputEdge] = []


class ELKInputChild(ELKInputData, total=False):
class ELKInputChild(ELKInputData):
"""Children of either `ELKInputData` or `ELKInputChild`."""

labels: cabc.MutableSequence[ELKInputLabel]
ports: cabc.MutableSequence[ELKInputPort]
labels: cabc.MutableSequence[ELKInputLabel] = []
ports: cabc.MutableSequence[ELKInputPort] = []

width: t.Union[int, float]
height: t.Union[int, float]
width: t.Union[int, float] = 0
height: t.Union[int, float] = 0


class ELKInputLabel(te.TypedDict, total=False):
class ELKInputLabel(BaseELKModel):
"""Label data that can be fed to ELK."""

text: te.Required[str]
layoutOptions: LayoutOptions
width: t.Union[int, float]
height: t.Union[int, float]
text: str
layoutOptions: LayoutOptions = {}
width: t.Union[int, float] = 0
height: t.Union[int, float] = 0


class ELKInputPort(t.TypedDict):
class ELKInputPort(BaseELKModel):
"""Connector data that can be fed to ELK."""

id: str
layoutOptions: LayoutOptions = {}

width: t.Union[int, float]
height: t.Union[int, float]

layoutOptions: te.NotRequired[cabc.MutableMapping[str, t.Any]]


class ELKInputEdge(te.TypedDict):
class ELKInputEdge(BaseELKModel):
"""Exchange data that can be fed to ELK."""

id: str
sources: cabc.MutableSequence[str]
targets: cabc.MutableSequence[str]
labels: te.NotRequired[cabc.MutableSequence[ELKInputLabel]]
labels: cabc.MutableSequence[ELKInputLabel] = []


class ELKPoint(t.TypedDict):
class ELKPoint(BaseELKModel):
"""Point data in ELK."""

x: t.Union[int, float]
y: t.Union[int, float]


class ELKSize(t.TypedDict):
class ELKSize(BaseELKModel):
"""Size data in ELK."""

width: t.Union[int, float]
height: t.Union[int, float]


class ELKOutputElement(t.TypedDict):
class ELKOutputElement(BaseELKModel):
"""Base class for all elements that comes out of ELK."""

id: str

style: dict[str, t.Any]
style: dict[str, t.Any] = {}


class ELKOutputData(ELKOutputElement):
"""Data that comes from ELK."""

type: t.Literal["graph"]
children: cabc.MutableSequence[ELKOutputChild] # type: ignore
children: cabc.MutableSequence[ELKOutputChild] = []


class ELKOutputNode(ELKOutputElement):
"""Node that comes out of ELK."""

type: t.Literal["node"]
children: cabc.MutableSequence[ELKOutputChild] # type: ignore
children: cabc.MutableSequence[ELKOutputChild] = []

position: ELKPoint
size: ELKSize
Expand All @@ -179,7 +188,7 @@ class ELKOutputJunction(ELKOutputElement):
"""Exchange-Junction that comes out of ELK."""

type: t.Literal["junction"]
children: cabc.MutableSequence[ELKOutputLabel]
children: cabc.MutableSequence[ELKOutputLabel] = []

position: ELKPoint
size: ELKSize
Expand All @@ -189,7 +198,7 @@ class ELKOutputPort(ELKOutputElement):
"""Port that comes out of ELK."""

type: t.Literal["port"]
children: cabc.MutableSequence[ELKOutputLabel]
children: cabc.MutableSequence[ELKOutputLabel] = []

position: ELKPoint
size: ELKSize
Expand All @@ -213,10 +222,10 @@ class ELKOutputEdge(ELKOutputElement):
sourceId: str
targetId: str
routingPoints: cabc.MutableSequence[ELKPoint]
children: cabc.MutableSequence[ELKOutputLabel]
children: cabc.MutableSequence[ELKOutputLabel] = []


ELKOutputChild = t.Union[ # type: ignore
ELKOutputChild = t.Union[
ELKOutputEdge,
ELKOutputJunction,
ELKOutputLabel,
Expand Down Expand Up @@ -344,15 +353,15 @@ def call_elkjs(elk_dict: ELKInputData) -> ELKOutputData:
executable=shutil.which("node"),
capture_output=True,
check=False,
input=json.dumps(elk_dict),
input=elk_dict.model_dump_json(),
text=True,
env={**os.environ, "NODE_PATH": str(NODE_HOME)},
)
if proc.returncode:
log.getChild("node").error("%s", proc.stderr)
raise NodeJSError("elk.js process failed")

return json.loads(proc.stdout)
return ELKOutputData.model_validate_json(proc.stdout)


def get_global_layered_layout_options() -> LayoutOptions:
Expand Down
18 changes: 9 additions & 9 deletions capellambse_context_diagrams/collectors/dataflow_view.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ def collector_portless(
)
made_edges: set[str] = set()
for act in activities:
data["children"].append(act_box := makers.make_box(act))
data.children.append(act_box := makers.make_box(act))
connections = list(portless.get_exchanges(act, filter=filter))

in_act: dict[str, oa.OperationalActivity] = {}
Expand All @@ -84,9 +84,9 @@ def collector_portless(
else:
in_act.setdefault(edge.target.uuid, edge.target)

act_box["height"] += (
makers.PORT_SIZE + 2 * makers.PORT_PADDING
) * max(len(in_act), len(out_act))
act_box.height += (makers.PORT_SIZE + 2 * makers.PORT_PADDING) * max(
len(in_act), len(out_act)
)

ex_datas: list[generic.ExchangeData] = []
for ex in connections:
Expand Down Expand Up @@ -124,7 +124,7 @@ def collector_default(
)
made_edges: set[str] = set()
for fnc in functions:
data["children"].append(fnc_box := makers.make_box(fnc))
data.children.append(fnc_box := makers.make_box(fnc))
_ports = default.port_collector(fnc, diagram.type)
connections = default.port_exchange_collector(_ports, filter=filter)
in_ports: dict[str, fa.FunctionPort] = {}
Expand All @@ -135,12 +135,12 @@ def collector_default(
else:
in_ports.setdefault(edge.target.uuid, edge.target)

fnc_box["ports"] = [
fnc_box.ports = [
makers.make_port(i.uuid) for i in (in_ports | out_ports).values()
]
fnc_box["height"] += (
makers.PORT_SIZE + 2 * makers.PORT_PADDING
) * max(len(in_ports), len(out_ports))
fnc_box.height += (makers.PORT_SIZE + 2 * makers.PORT_PADDING) * max(
len(in_ports), len(out_ports)
)

ex_datas: list[generic.ExchangeData] = []
for ex in edges:
Expand Down
64 changes: 30 additions & 34 deletions capellambse_context_diagrams/collectors/default.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,9 +40,9 @@ def collector(
)
data = generic.collector(diagram, no_symbol=True)
ports = port_collector(diagram.target, diagram.type)
centerbox = data["children"][0]
centerbox = data.children[0]
connections = port_exchange_collector(ports)
centerbox["ports"] = [
centerbox.ports = [
makers.make_port(uuid) for uuid, edges in connections.items() if edges
]
ex_datas: list[generic.ExchangeData] = []
Expand All @@ -53,9 +53,9 @@ def collector(
if is_hierarchical := exchanges.is_hierarchical(ex, centerbox):
if not diagram.display_parent_relation:
continue
centerbox["labels"][0][
"layoutOptions"
] = makers.DEFAULT_LABEL_LAYOUT_OPTIONS
centerbox.labels[0].layoutOptions = (
makers.DEFAULT_LABEL_LAYOUT_OPTIONS
)
elkdata: _elkjs.ELKInputData = centerbox
else:
elkdata = data
Expand All @@ -67,9 +67,9 @@ def collector(
ex_datas.append(ex_data)
except AttributeError:
continue
global_boxes = {centerbox["id"]: centerbox}
made_boxes = {centerbox["id"]: centerbox}
boxes_to_delete = {centerbox["id"]}
global_boxes = {centerbox.id: centerbox}
made_boxes = {centerbox.id: centerbox}
boxes_to_delete = {centerbox.id}

def _make_box_and_update_globals(
obj: t.Any,
Expand All @@ -90,8 +90,8 @@ def _make_owner_box(current: t.Any) -> t.Any:
no_symbol=diagram.display_symbols_as_boxes,
layout_options=makers.DEFAULT_LABEL_LAYOUT_OPTIONS,
)
for box in (children := parent_box.setdefault("children", [])):
if box["id"] == current.uuid:
for box in (children := parent_box.children):
if box.id == current.uuid:
box = global_boxes.get(current.uuid, current)
break
else:
Expand All @@ -107,8 +107,8 @@ def _make_owner_box(current: t.Any) -> t.Any:
no_symbol=diagram.display_symbols_as_boxes,
layout_options=makers.DEFAULT_LABEL_LAYOUT_OPTIONS,
)
box["children"] = [centerbox]
del data["children"][0]
box.children = [centerbox]
del data.children[0]
except AttributeError:
pass
diagram_target_owners = generic.get_all_owners(diagram.target)
Expand All @@ -128,17 +128,15 @@ def _make_owner_box(current: t.Any) -> t.Any:
if box := global_boxes.get(child.uuid): # type: ignore[assignment]
if box is centerbox:
continue
box.setdefault("ports", []).extend(
[makers.make_port(j.uuid) for j in local_ports]
)
box["height"] += height
box.ports.extend([makers.make_port(j.uuid) for j in local_ports])
box.height += height
else:
box = _make_box_and_update_globals(
child,
height=height,
no_symbol=diagram.display_symbols_as_boxes,
)
box["ports"] = [makers.make_port(j.uuid) for j in local_ports]
box.ports = [makers.make_port(j.uuid) for j in local_ports]

if diagram.display_parent_relation:
current = child
Expand Down Expand Up @@ -170,17 +168,15 @@ def _make_owner_box(current: t.Any) -> t.Any:

for uuid in boxes_to_delete:
del global_boxes[uuid]
data["children"].extend(global_boxes.values())
data.children.extend(global_boxes.values())
if diagram.display_parent_relation:
owner_boxes: dict[str, _elkjs.ELKInputChild] = {
uuid: box
for uuid, box in made_boxes.items()
if box.get("children")
uuid: box for uuid, box in made_boxes.items() if box.children
}
generic.move_parent_boxes_to_owner(owner_boxes, diagram.target, data)
generic.move_edges(owner_boxes, edges, data)

centerbox["height"] = max(centerbox["height"], *stack_heights.values())
centerbox.height = max(centerbox.height, *stack_heights.values())
derivator = DERIVATORS.get(type(diagram.target))
if diagram.display_derived_interfaces and derivator is not None:
derivator(diagram, data, made_boxes[diagram.target.uuid])
Expand Down Expand Up @@ -317,7 +313,7 @@ def derive_from_functions(
for fnc in diagram.target.allocated_functions:
ports.extend(port_collector(fnc, diagram.type))

context_box_ids = {child["id"] for child in data["children"]}
context_box_ids = {child.id for child in data.children}
components: dict[str, cs.Component] = {}
for port in ports:
for fex in port.exchanges:
Expand Down Expand Up @@ -349,24 +345,24 @@ def derive_from_functions(
no_symbol=diagram.display_symbols_as_boxes,
)
class_ = type(derived_comp).__name__
box["id"] = f"{STYLECLASS_PREFIX}-{class_}:{uuid}"
data["children"].append(box)
box.id = f"{STYLECLASS_PREFIX}-{class_}:{uuid}"
data.children.append(box)
source_id = f"{STYLECLASS_PREFIX}-CP_INOUT:{i}"
target_id = f"{STYLECLASS_PREFIX}-CP_INOUT:{-i}"
box.setdefault("ports", []).append(makers.make_port(source_id))
centerbox.setdefault("ports", []).append(makers.make_port(target_id))
box.ports.append(makers.make_port(source_id))
centerbox.ports.append(makers.make_port(target_id))
if i % 2 == 0:
source_id, target_id = target_id, source_id

data["edges"].append(
{
"id": f"{STYLECLASS_PREFIX}-ComponentExchange:{i}",
"sources": [source_id],
"targets": [target_id],
}
data.edges.append(
_elkjs.ELKInputEdge(
id=f"{STYLECLASS_PREFIX}-ComponentExchange:{i}",
sources=[source_id],
targets=[target_id],
)
)

data["children"][0]["height"] += (
data.children[0].height += (
makers.PORT_PADDING
+ (makers.PORT_SIZE + makers.PORT_PADDING) * len(components) // 2
)
Expand Down
Loading

0 comments on commit 61cdead

Please sign in to comment.