make use of pydantic types for interface data

napalm-automation-community · Aug 23, 2024 · 2865c38 · 2865c38
1 parent 6c53bef
commit 2865c38
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 12 deletions.
diff --git a/napalm_srlinux/srlinux.py b/napalm_srlinux/srlinux.py
@@ -37,6 +37,8 @@
     ReplaceConfigException,
 )
 
+from napalm_srlinux import types
+
 
 class NokiaSRLinuxDriver(NetworkDriver):
     """Napalm driver for Nokia SR Linux."""
@@ -273,7 +275,7 @@ def get_facts(self):
 
         raise NotImplementedError
 
-    def get_interfaces(self) -> dict:
+    def get_interfaces(self) -> dict[str, types.Interface]:
         """
         Returns a dictionary of dictionaries.
         The keys for the first dictionary will be the interfaces in the devices.
@@ -303,21 +305,19 @@ def get_interfaces(self) -> dict:
         system_data = json_data[1]
 
         for interface in interfaces_json.get("srl_nokia-interfaces:interface"):
-            interfaces[interface.get("name")] = {
-                "is_up": True if interface.get("oper-state") == "up" else False,
-                "is_enabled": True
-                if interface.get("admin-state") == "enable"
-                else False,
-                "description": interface.get("description"),
-                "last-flapped": NokiaSRLinuxDriver._calculate_time_since(
+            interfaces[interface.get("name")] = types.Interface(
+                is_up=interface.get("oper-state") == "up",
+                is_enabled=interface.get("admin-state") == "enable",
+                description=interface.get("description"),
+                last_flapped=NokiaSRLinuxDriver._calculate_time_since(
                     system_data.get("current-datetime"), interface.get("last-change")
                 ),
-                "speed": NokiaSRLinuxDriver._port_speed_to_mbits(
+                speed=NokiaSRLinuxDriver._port_speed_to_mbits(
                     interface.get("ethernet", {}).get("port-speed")
                 ),
-                "mtu": interface.get("mtu"),
-                "mac_address": interface.get("ethernet", {}).get("hw-mac-address"),
-            }
+                mtu=interface.get("mtu"),
+                mac_address=interface.get("ethernet", {}).get("hw-mac-address"),
+            )
 
         return interfaces
 

diff --git a/napalm_srlinux/types.py b/napalm_srlinux/types.py
@@ -0,0 +1,11 @@
+from pydantic import BaseModel
+
+
+class Interface(BaseModel):
+    is_up: bool
+    is_enabled: bool
+    description: str | None
+    last_flapped: float
+    speed: float | None
+    mtu: int | None
+    mac_address: str | None