Merge pull request #71 from digitalocean/custom-field-serialization

Custom field serialization

Author: Zach Moody

pynetbox/lib/response.py

@@ -19,9 +19,9 @@
 from pynetbox.lib.query import Request
 
 
-def _get_return(lookup, return_fields=['id', 'value', 'nested_return']):
+def get_return(lookup, return_fields=None):
 
-    for i in return_fields:
+    for i in return_fields or ['id', 'value', 'nested_return']:
         if isinstance(lookup, dict) and lookup.get(i):
             return lookup[i]
         else:
@@ -33,6 +33,13 @@ def _get_return(lookup, return_fields=['id', 'value', 'nested_return']):
         return lookup
 
 
+def flatten_custom(custom_dict):
+    return {
+        k: v if not isinstance(v, dict) else v['value']
+        for k, v in custom_dict.items()
+    }
+
+
 class Record(object):
     """Create python objects from netbox API responses.
 
@@ -114,7 +121,7 @@ def _add_cache(self, item):
             self._full_cache.append((key, dict(value)))
         else:
             self._full_cache.append((key, value))
-        self._index_cache.append((key, _get_return(value)))
+        self._index_cache.append((key, get_return(value)))
 
     def _parse_values(self, values):
         """ Parses values init arg.
@@ -148,12 +155,11 @@ def _compare(self):
         init_vals = dict(self._index_cache)
         for i in dict(self):
             current_val = init_vals.get(i)
-            if i != 'custom_fields':
-                if isinstance(current_val, dict):
-                    init_dict.update({i: _get_return(current_val)})
-                else:
-                    init_dict.update({i: _get_return(current_val)})
-            init_dict.update({i: current_val})
+            if i == 'custom_fields':
+                init_dict[i] = flatten_custom(current_val)
+            else:
+                init_dict[i] = get_return(current_val)
+
         if init_dict == self.serialize():
             return True
         return False
@@ -193,18 +199,24 @@ def serialize(self, nested=False):
         :returns: dict of values the NetBox API is expecting.
         """
         if nested:
-            return _get_return(self)
+            return get_return(self)
 
         ret = {}
         for i in dict(self):
             current_val = getattr(self, i)
-            if isinstance(current_val, Record):
-                current_val = getattr(current_val, 'serialize')(nested=True)
+            if i == 'custom_fields':
+                ret[i] = flatten_custom(current_val)
+            else:
+                if isinstance(current_val, Record):
+                    current_val = getattr(
+                        current_val,
+                        'serialize'
+                    )(nested=True)
 
-            if isinstance(current_val, netaddr.ip.IPNetwork):
-                current_val = str(current_val)
+                if isinstance(current_val, netaddr.ip.IPNetwork):
+                    current_val = str(current_val)
 
-            ret.update({i: current_val})
+                ret[i] = current_val
         return ret
 
     def save(self):

tests/fixtures/dcim/site.json

@@ -13,7 +13,11 @@
     "contact_email": "",
     "comments": "",
     "custom_fields": {
-        "test_custom": "Hello"
+        "test_custom": "Hello",
+        "test_selection": {
+            "value": 2,
+            "label": "second"
+        }
     },
     "count_prefixes": 2,
     "count_vlans": 1,

tests/test_dcim.py

@@ -282,7 +282,7 @@ class SiteTestCase(unittest.TestCase, GenericTest):
         'pynetbox.lib.query.requests.get',
         return_value=Response(fixture='dcim/site.json')
     )
-    def test_modify(self, mock):
+    def test_modify_custom(self, mock):
         '''Test modifying a custom field.
         '''
         ret = getattr(nb, self.name).get(1)
@@ -293,6 +293,23 @@ def test_modify(self, mock):
             ret.custom_fields['test_custom'], 'Testing'
         )
 
+    @patch(
+        'pynetbox.lib.query.requests.get',
+        return_value=Response(fixture='dcim/site.json')
+    )
+    def test_custom_selection_serializer(self, _):
+        '''Tests serializer with custom selection fields.
+        '''
+        ret = getattr(nb, self.name).get(1)
+        ret.custom_fields['test_custom'] = "Testing"
+        test = ret.serialize()
+        from pprint import pprint
+        pprint(test)
+        self.assertEqual(
+            test['custom_fields']['test_selection'],
+            2
+        )
+
     @patch(
         'pynetbox.lib.query.requests.post',
         return_value=Response(fixture='dcim/site.json')