Fix lint issues with HMC7044 and draw classes

Signed-off-by: Travis F. Collins <[email protected]>

adijif/clocks/hmc7044.py

@@ -164,8 +164,8 @@ def vxco_doubler(self, value: Union[int, List[int]]) -> None:
         self._check_in_range(value, self.vxco_doubler_available, "vxco_doubler")
         self._vxco_doubler = value
 
-    def _init_diagram(self):
-        """Initialize diagram for HMC7044 alone"""
+    def _init_diagram(self) -> None:
+        """Initialize diagram for HMC7044 alone."""
         self.ic_diagram_node = None
         self._diagram_output_dividers = []
 
@@ -223,11 +223,14 @@ def _init_diagram(self):
         #     self.ic_diagram_node.add_connection({"from": out_dividers, "to": div})
         #     # root.add_connection({"from": vco, "to": div})
 
-    def _update_diagram(self, config: Dict):
+    def _update_diagram(self, config: Dict) -> None:
         """Update diagram with configuration.
 
         Args:
             config (Dict): Configuration dictionary
+
+        Raises:
+            Exception: If key is not D followed by a number
         """
         # Add output dividers
         keys = config.keys()
@@ -244,8 +247,15 @@ def _update_diagram(self, config: Dict):
                     f"Unknown key {key}. Must be of for DX where X is a number"
                 )
 
-    def draw(self):
-        """Draw diagram in d2 language for IC alone with reference clock."""
+    def draw(self) -> str:
+        """Draw diagram in d2 language for IC alone with reference clock.
+
+        Returns:
+            str: Diagram in d2 language
+
+        Raises:
+            Exception: If no solution is saved
+        """
         if not self._saved_solution:
             raise Exception("No solution to draw. Must call solve first.")
         lo = Layout("HMC7044 Example")
@@ -267,12 +277,12 @@ def draw(self):
         node.value = str(self._saved_solution["n2"])
 
         # Update VCXO Doubler to R2
-        con = self.ic_diagram_node.get_connection("VCXO Doubler", "R2")
+        # con = self.ic_diagram_node.get_connection("VCXO Doubler", "R2")
         rate = self._saved_solution["vcxo_doubler"] * self._saved_solution["vcxo"]
         self.ic_diagram_node.update_connection("VCXO Doubler", "R2", rate)
 
         # Update R2 to PFD
-        con = self.ic_diagram_node.get_connection("R2", "PFD")
+        # con = self.ic_diagram_node.get_connection("R2", "PFD")
         rate = (
             self._saved_solution["vcxo"]
             * self._saved_solution["vcxo_doubler"]
@@ -281,7 +291,7 @@ def draw(self):
         self.ic_diagram_node.update_connection("R2", "PFD", rate)
 
         # Update VCO
-        con = self.ic_diagram_node.get_connection("VCO", "Output Dividers")
+        # con = self.ic_diagram_node.get_connection("VCO", "Output Dividers")
         self.ic_diagram_node.update_connection(
             "VCO", "Output Dividers", self._saved_solution["vco"]
         )

adijif/draw.py

@@ -1,9 +1,21 @@
-"""Diagraming functions for different componets."""
+"""Diagraming functions for different components."""
+from __future__ import annotations
+
 import os
+from typing import Union
 
 
 class Node:
-    def __init__(self, name, ntype=None, parent=None):
+    """Node model for diagraming which can have children and connections."""
+
+    def __init__(self, name: str, ntype: str = None, parent: Node = None) -> None:
+        """Initialize node with name, type and parent node.
+
+        Args:
+            name (str): Name of the node.
+            ntype (str): Type of the node.
+            parent (Node): Parent node of the node.
+        """
         self.name = name
         self.parent = parent
         self.ntype = ntype
@@ -13,69 +25,129 @@ def __init__(self, name, ntype=None, parent=None):
         self.use_unit_conversion_for_rate = True
         self._value = None
 
-    def __repr__(self):
+    def __repr__(self) -> str:
+        """Get string representation of the node.
+
+        Returns:
+            str: String representation of the node.
+        """
         return f"Node({self.name})"
 
     @property
-    def value(self):
+    def value(self) -> str:
+        """Get value of the node.
+
+        Returns:
+            str: Value of the node.
+        """
         return f"{self.name} = {self._value}" if self._value else None
 
     @value.setter
-    def value(self, value):
+    def value(self, value: Union(int, float, str)) -> None:
+        """Set value of the node."""
         self._value = value
 
-    def add_connection(self, connection: dict):
+    def add_connection(self, connection: dict) -> None:
+        """Add connection between this node and another node.
+
+        Args:
+            connection (dict): Connection dictionary with keys "from", "to"
+                and optionally "rate".
+        """
         if "rate" in connection and self.use_unit_conversion_for_rate:
             units = ["Hz", "kHz", "MHz", "GHz"]
             rate = connection["rate"]
             # Convert to unit based on rate scale
-            for i, unit in enumerate(units):
+            for unit in units:
                 if rate < 1000:
+                    connection["rate"] = f"{rate:.2f} {unit}"
                     break
                 rate /= 1000
-            connection["rate"] = f"{rate:.2f} {unit}"
 
         self.connections.append(connection)
 
-    def get_connection(self, from_s, to):
+    def get_connection(self, from_s: str, to: str) -> dict:
+        """Get connection between this node and another node.
+
+        Args:
+            from_s (str): Name of the node from which connection originates.
+            to (str): Name of the node to which connection goes.
+
+        Returns:
+            dict: Connection dictionary with keys "from", "to" and optionally "rate".
+
+        Raises:
+            ValueError: If connection not found.
+        """
         for conn in self.connections:
             if conn["from"].name == from_s and conn["to"].name == to:
                 return conn
         raise ValueError(f"Connection from {from_s} to {to} not found.")
 
-    def update_connection(self, from_s, to, rate):
+    def update_connection(self, from_s: str, to: str, rate: Union(int, float)) -> None:
+        """Update connection rate between this node and another node.
+
+        Args:
+            from_s (str): Name of the node from which connection originates.
+            to (str): Name of the node to which connection goes.
+            rate (float): Rate of the connection.
+
+        Raises:
+            ValueError: If connection not found.
+        """
         for conn in self.connections:
             if conn["from"].name == from_s and conn["to"].name == to:
                 units = ["Hz", "kHz", "MHz", "GHz"]
                 # Convert to unit based on rate scale
-                for i, unit in enumerate(units):
+                for unit in units:
                     if rate < 1000:
-                        break
+                        conn["rate"] = f"{rate:.2f} {unit}"
+                        return
                     rate /= 1000
-                conn["rate"] = f"{rate:.2f} {unit}"
-                return
 
         raise ValueError(f"Connection from {from_s} to {to} not found.")
 
-    def add_child(self, child):
+    def add_child(self, child: Node) -> None:
+        """Add child node to this node.
+
+        Args:
+            child (Node): Child node to add.
+        """
         if not isinstance(child, list):
             child = [child]
         for c in child:
             c.parent = self
             self.children.append(c)
 
-    def get_child(self, name: str):
+    def get_child(self, name: str) -> Node:
+        """Get child node by name.
+
+        Args:
+            name (str): Name of the child node.
+
+        Returns:
+            Node: Child node with the given name.
+
+        Raises:
+            ValueError: If child node not found.
+        """
         for child in self.children:
             if child.name == name:
                 return child
         raise ValueError(f"Child with name {name} not found.")
 
 
 class Layout:
+    """Layout model for diagraming which contains nodes and connections."""
 
     si = "    "
 
-    def __init__(self, name: str):
+    def __init__(self, name: str) -> None:
+        """Initialize layout with name.
+
+        Args:
+            name (str): Name of the layout.
+        """
         self.name = name
         self.nodes = []
         self.connections = []
@@ -84,27 +156,49 @@ def __init__(self, name: str):
         self.output_image_filename = "clocks.png"
         self.layout_engine = "elk"
 
-    def add_node(self, node: Node):
+    def add_node(self, node: Node) -> None:
+        """Add node to the layout.
+
+        Args:
+            node (Node): Node to add to the layout.
+        """
         self.nodes.append(node)
 
-    def add_connection(self, connection: dict):
+    def add_connection(self, connection: dict) -> None:
+        """Add connection between two nodes.
+
+        Args:
+            connection (dict): Connection dictionary with keys "from", "to"
+                and optionally "rate".
+        """
         if "rate" in connection and self.use_unit_conversion_for_rate:
             units = ["Hz", "kHz", "MHz", "GHz"]
             rate = connection["rate"]
             # Convert to unit based on rate scale
-            for i, unit in enumerate(units):
+            for unit in units:
                 if rate < 1000:
+                    connection["rate"] = f"{rate:.2f} {unit}"
                     break
                 rate /= 1000
-            connection["rate"] = f"{rate:.2f} {unit}"
         self.connections.append(connection)
 
-    def draw(self):
-        """Draw diagram in d2 language."""
+    def draw(self) -> str:
+        """Draw diagram in d2 language.
 
+        Returns:
+            str: Path to the output image file.
+        """
         diag = "direction: right\n\n"
 
-        def get_parents_names(node):
+        def get_parents_names(node: Node) -> str:
+            """Get names of all parent nodes of the given node.
+
+            Args:
+                node (Node): Node for which to get parent names.
+
+            Returns:
+                str: Names of all parent nodes of the given node.
+            """
             parents = []
             while node.parent:
                 parents.append(node.parent.name)
@@ -113,7 +207,16 @@ def get_parents_names(node):
                 return ""
             return ".".join(parents[::-1]) + "."
 
-        def draw_subnodes(node, spacing="    "):
+        def draw_subnodes(node: Node, spacing: str = "    ") -> str:
+            """Draw subnodes of the given node.
+
+            Args:
+                node (Node): Node for which to draw subnodes.
+                spacing (str): Spacing for indentation.
+
+            Returns:
+                str: Subnodes of the given node.
+            """
             diag = " {\n"
             for child in node.children:
                 if child.value:
@@ -150,7 +253,8 @@ def draw_subnodes(node, spacing="    "):
             from_p_name = get_parents_names(connection["from"])
             to_p_name = get_parents_names(connection["to"])
             label = f"{connection['rate']}" if "rate" in connection else None
-            diag += f"{from_p_name}{connection['from'].name} -> {to_p_name}{connection['to'].name}"
+            diag += f"{from_p_name}{connection['from'].name} ->"
+            diag += f" {to_p_name}{connection['to'].name}"
             diag += ": " + label if label else ""
             diag += "\n"
 
@@ -159,15 +263,16 @@ def draw_subnodes(node, spacing="    "):
                 from_p_name = get_parents_names(connection["from"])
                 to_p_name = get_parents_names(connection["to"])
                 label = f"{connection['rate']}" if "rate" in connection else ""
-                diag += f"{from_p_name}{connection['from'].name} -> {to_p_name}{connection['to'].name}"
+                diag += f"{from_p_name}{connection['from'].name} -> "
+                diag += f"{to_p_name}{connection['to'].name}"
                 diag += ": " + label if label else ""
                 diag += "\n"
 
         with open(self.output_filename, "w") as f:
             f.write(diag)
 
-        os.system(
-            f"d2 -l {self.layout_engine} {self.output_filename} {self.output_image_filename}"
-        )
+        cmd = f"d2 -l {self.layout_engine} {self.output_filename} "
+        cmd += f"{self.output_image_filename}"
+        os.system(cmd)  # noqa: S605
 
         return self.output_image_filename