add tags to load

Author: chun-tseng

ns1/ipam.py

@@ -46,19 +46,21 @@ class OptiondefException(Exception):
 
 
 class Network(object):
-    def __init__(self, config, name=None, id=None):
+    def __init__(self, config, name=None, id=None, tags=None):
         """
         Create a new high level Network object
 
         :param ns1.config.Config config: config object
         :param str name: network name
         :param int id: id of an existing Network
+        :param dict tags: tags of the reservation
         """
         self._rest = Networks(config)
         self.config = config
         self.name = name
         self.id = id
         self.report = {}
+        self.tags = tags
         self.data = None
 
     def __repr__(self):
@@ -88,6 +90,9 @@ def success(result, *args):
             self.name = result["name"]
             self.report = self._rest.report(self.id)
 
+            if "tags" in result:
+                self.tags = result["tags"]
+
             if callback:
                 return callback(self)
             else:
@@ -161,7 +166,7 @@ def success(result, *args):
         )
 
     def new_address(
-        self, prefix, status, callback=None, errback=None, **kwargs
+            self, prefix, status, callback=None, errback=None, **kwargs
     ):
         """
         Create a new address space in this Network
@@ -179,13 +184,14 @@ def new_address(
 
 class Address(object):
     def __init__(
-        self,
-        config,
-        prefix=None,
-        status=None,
-        network=None,
-        scope_group=None,
-        id=None,
+            self,
+            config,
+            prefix=None,
+            status=None,
+            network=None,
+            scope_group=None,
+            id=None,
+            tags=None
     ):
         """
         Create a new high level Address object
@@ -195,6 +201,7 @@ def __init__(
         :param str status: planned, assigned
         :param Network network: Network Object the address will be part of
         :param Scopegroup scope_group: Scopegroup Object that will be associated with the address
+        :param dict tags: tags of the reservation
         """
         self._rest = Addresses(config)
         self.config = config
@@ -206,6 +213,7 @@ def __init__(
         self.children = []
         self.report = {}
         self.data = None
+        self.tags = tags
 
     def __repr__(self):
         return "<Address address=%s>" % self.prefix
@@ -237,6 +245,10 @@ def success(result, *args):
             # self.scope_group = Scopegroup(config=self.config, id=result['scope_group_id']) NYI
             self.report = self._rest.report(self.id)
             children = self._rest.retrieve_children(self.id)
+
+            if "tags" in result:
+                self.tags = result["tags"]
+
             self.children = [
                 Address(self.config, id=child["id"])
                 for child in children
@@ -255,9 +267,9 @@ def success(result, *args):
 
         if self.id is None:
             if (
-                self.prefix is None
-                or self.status is None
-                or self.network is None
+                    self.prefix is None
+                    or self.status is None
+                    or self.network is None
             ):
                 raise AddressException(
                     "Must at least specify an id or prefix, status, and network"
@@ -269,8 +281,8 @@ def success(result, *args):
                         address
                         for address in self._rest.list()
                         if address["prefix"] == self.prefix
-                        and address["status"] == self.status
-                        and address["network_id"] == network_id
+                           and address["status"] == self.status
+                           and address["network_id"] == network_id
                     ][0]["id"]
                 except IndexError:
                     raise AddressException(
@@ -326,7 +338,7 @@ def success(result, *args):
         )
 
     def reserve(
-        self, scopegroup_id, mac, options=None, callback=None, errback=None
+            self, scopegroup_id, mac, options=None, callback=None, errback=None
     ):
         """
         Add scope group reservation. Pass a single Address object and a MAC address as a string
@@ -390,14 +402,15 @@ def success(result, *args):
 
 
 class Scopegroup(object):
-    def __init__(self, config, name=None, service_def_id=None, id=None):
+    def __init__(self, config, name=None, service_def_id=None, id=None, tags=None):
         """
         Create a new high level Scopegroup object
 
         :param ns1.config.Config config: config object
         :param str name: Name of the scope group
         :param int service_group_id: id of the service group the scope group is associated with
         :param int id: id of the scope group
+        :param dict tags: tags of the reservation
         """
         self._rest = Scopegroups(config)
         self.config = config
@@ -407,6 +420,7 @@ def __init__(self, config, name=None, service_def_id=None, id=None):
         self.name = name
         self.dhcp_service_id = service_def_id
         self.data = None
+        self.tags = tags
 
     def __repr__(self):
         return "<Scopegroup scope_group=%s>" % self.name
@@ -435,6 +449,9 @@ def success(result, *args):
             self.dhcp4 = result["dhcpv4"]
             self.dhcp6 = result["dhcpv6"]
             self.dhcp_service_id = result.get("dhcp_service_id")
+            
+            if "tags" in result:
+                self.tags = result["tags"]
 
             if callback:
                 return callback(self)
@@ -522,7 +539,7 @@ def success(result, *args):
         )
 
     def reserve(
-        self, address_id, mac, options=None, callback=None, errback=None
+            self, address_id, mac, options=None, callback=None, errback=None
     ):
         """
         :param int address_id: id of the Address to reserve
@@ -576,14 +593,14 @@ def scopes(self, callback=None, errback=None):
 
 class Reservation(object):
     def __init__(
-        self,
-        config,
-        scopegroup_id,
-        address_id,
-        tags=None,
-        reservation_id=None,
-        options=None,
-        mac=None,
+            self,
+            config,
+            scopegroup_id,
+            address_id,
+            reservation_id=None,
+            options=None,
+            mac=None,
+            tags=None
     ):
         """
         Create a new high level Reservation object
@@ -595,15 +612,16 @@ def __init__(
         :param int reservation_id: id of the reservation
         :param list options: dhcp options of the reservation
         :param str mac: mac address of the reservation
+        :param dict tags: tags of the reservation
         """
         self._rest = Reservations(config)
         self.config = config
         self.id = reservation_id
         self.scopegroup_id = scopegroup_id
-        self.tags = tags
         self.address_id = address_id
         self.mac = mac
         self.data = None
+        self.tags = tags
 
         if options is None:
             options = DHCPOptions("dhcpv4", {})
@@ -698,7 +716,7 @@ def success(result, *args):
         )
 
     def update(
-        self, options, callback=None, errback=None, parent=True, **kwargs
+            self, options, callback=None, errback=None, parent=True, **kwargs
     ):
         """
         Update reservation configuration. Pass a list of keywords and their values to
@@ -825,7 +843,7 @@ def success(result, *args):
 
 class Scope(object):
     def __init__(
-        self, config, scopegroup_id, address_id, scope_id=None, options=None
+            self, config, scopegroup_id, address_id, scope_id=None, options=None, tags=None
     ):
         """
         Create a new high level Scope object
@@ -835,12 +853,14 @@ def __init__(
         :param int address_id: id of the address the scope is associated with
         :param int scope_id: id of the scope
         :param DHCPOptions options: DHCPOptions object that contains the settings for the scope
+        :param dict tags: tags of the reservation
         """
         self._rest = Scopes(config)
         self.config = config
         self.scopegroup_id = scopegroup_id
         self.address_id = address_id
         self.id = scope_id
+        self.tags = tags
 
         if options is None:
             options = DHCPOptions("dhcpv4", {})
@@ -883,6 +903,9 @@ def success(result, *args):
             self.address_id = result["address_id"]
             self.options = result["options"]
 
+            if "tags" in result:
+                self.tags = result["tags"]
+
             if callback:
                 return callback(self)
             else:
@@ -930,7 +953,7 @@ def success(result, *args):
         )
 
     def update(
-        self, address_id, options, callback=None, errback=None, **kwargs
+            self, address_id, options, callback=None, errback=None, **kwargs
     ):
         """
         Update Scope configuration. Pass a list of keywords and their values to
@@ -984,14 +1007,14 @@ def reload(self, callback=None, errback=None):
         return self.load(reload=True, callback=callback, errback=errback)
 
     def load(
-        self,
-        scope_group_id=None,
-        scope_id=None,
-        limit=None,
-        offset=None,
-        callback=None,
-        errback=None,
-        reload=False,
+            self,
+            scope_group_id=None,
+            scope_id=None,
+            limit=None,
+            offset=None,
+            callback=None,
+            errback=None,
+            reload=False,
     ):
         """
         Load Lease data from the API.