diff --git a/netutils/ip.py b/netutils/ip.py index 267e7329..7b2df0e2 100644 --- a/netutils/ip.py +++ b/netutils/ip.py @@ -2,7 +2,7 @@ import ipaddress import typing as t -from operator import attrgetter +from operator import attrgetter, methodcaller from netutils.constants import IPV4_MASKS, IPV6_MASKS @@ -60,12 +60,13 @@ def ipaddress_interface(ip: str, attr: str) -> t.Any: return retrieved_method -def ipaddress_network(ip: str, attr: str) -> t.Any: +def ipaddress_network(ip: str, attr: str, **kwargs: t.Any) -> t.Any: """Convenience function primarily built to expose ipaddress.ip_network to Jinja. Args: ip: IP network str compliant with ipaddress.ip_network inputs. attr: An attribute in string dotted format. + kwargs: Keyword arguments to pass along to the given method of ipaddress.ip_network. Returns: Returns the value provided by the ipaddress.ip_network attribute provided. @@ -76,9 +77,14 @@ def ipaddress_network(ip: str, attr: str) -> t.Any: 4 >>> ipaddress_network('10.1.1.0/24', '__str__') '10.1.1.0/24' - >>> + >>> list(ipaddress_network('192.168.1.0/28', 'subnets', new_prefix=30)) + [IPv4Network('192.168.1.0/30'), IPv4Network('192.168.1.4/30'), IPv4Network('192.168.1.8/30'), IPv4Network('192.168.1.12/30')] """ - retriever = attrgetter(attr) + retriever: t.Callable[[t.Union[ipaddress.IPv4Network, ipaddress.IPv6Network]], t.Any] + if kwargs: + retriever = methodcaller(attr, **kwargs) + else: + retriever = attrgetter(attr) retrieved_method = retriever(ipaddress.ip_network(ip)) if callable(retrieved_method): return retrieved_method() diff --git a/tests/unit/test_ip.py b/tests/unit/test_ip.py index b76447e0..28ec4d22 100644 --- a/tests/unit/test_ip.py +++ b/tests/unit/test_ip.py @@ -43,6 +43,17 @@ }, ] +IP_NETWORK_WITH_KWARGS = [ + { + "sent": {"ip": "10.1.1.0/28", "attr": "subnets", "new_prefix": 30}, + "received": "[IPv4Network('10.1.1.0/30'), IPv4Network('10.1.1.4/30'), IPv4Network('10.1.1.8/30'), IPv4Network('10.1.1.12/30')]", + }, + { + "sent": {"ip": "10.1.1.0/28", "attr": "subnets"}, + "received": "[IPv4Network('10.1.1.0/29'), IPv4Network('10.1.1.8/29')]", + }, +] + IP_NETWORK = [ { "sent": {"ip": "10.1.1.0/24", "attr": "hostmask.__str__"}, @@ -658,6 +669,11 @@ def test_ipaddress_network(data): assert ip.ipaddress_network(**data["sent"]) == data["received"] +@pytest.mark.parametrize("data", IP_NETWORK_WITH_KWARGS) +def test_ipaddress_network_with_kwargs(data): + assert str(list(ip.ipaddress_network(**data["sent"]))) == data["received"] + + @pytest.mark.parametrize("data", IS_CLASSFUL) def test_is_classful(data): assert ip.is_classful(**data["sent"]) == data["received"]