chore: change typing

canonical · Jan 26, 2024 · 4405ec5 · 4405ec5
1 parent bcc9aa5
commit 4405ec5
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/src-docs/firewall.py.md b/src-docs/firewall.py.md
@@ -12,7 +12,7 @@ The runner firewall manager.
 ## <kbd>class</kbd> `Firewall`
 Represent a firewall and provides methods to refresh its configuration. 
 
-<a href="../src/firewall.py#L50"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
+<a href="../src/firewall.py#L52"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
 
 ### <kbd>function</kbd> `__init__`
 
@@ -33,7 +33,7 @@ Initialize a new Firewall instance.
 
 ---
 
-<a href="../src/firewall.py#L58"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
+<a href="../src/firewall.py#L60"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
 
 ### <kbd>function</kbd> `get_host_ip`
 
@@ -50,7 +50,7 @@ Get the host IP address for the corresponding LXD network.
 
 ---
 
-<a href="../src/firewall.py#L96"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
+<a href="../src/firewall.py#L98"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
 
 ### <kbd>function</kbd> `refresh_firewall`
 
@@ -86,7 +86,7 @@ Represent an entry in the firewall.
 
 ---
 
-<a href="../src/firewall.py#L25"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
+<a href="../src/firewall.py#L27"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
 
 ### <kbd>classmethod</kbd> `decode`
 

diff --git a/src/firewall.py b/src/firewall.py
@@ -11,6 +11,8 @@
 
 from utilities import execute_command
 
+NetworkT = typing.TypeVar("NetworkT", ipaddress.IPv4Network, ipaddress.IPv6Network)
+
 
 @dataclasses.dataclass
 class FirewallEntry:
@@ -68,9 +70,9 @@ def get_host_ip(self) -> str:
 
     def _exclude_network(
         self,
-        networks: typing.List[ipaddress.IPv4Network],
-        excludes: typing.List[ipaddress.IPv4Network],
-    ) -> typing.Set[ipaddress.IPv4Network]:
+        networks: typing.List[NetworkT],
+        excludes: typing.List[NetworkT],
+    ) -> typing.Set[NetworkT]:
         """Excludes the network segment from a pool of networks.
 
         Args:
@@ -83,7 +85,7 @@ def _exclude_network(
         total_excluded_networks = set(networks)
 
         for exclude in excludes:
-            scoped_excluded_networks: set[ipaddress.IPv4Network] = set()
+            scoped_excluded_networks: set[NetworkT] = set()
             for net in total_excluded_networks:
                 if net.overlaps(exclude):
                     scoped_excluded_networks.update(net.address_exclude(exclude))