Properly use the JSONEncoders

pull/3/head
Raphaël Vinot 2017-11-01 11:20:09 -07:00
parent 38b6b72cb7
commit 12f5302eb3
3 changed files with 20 additions and 15 deletions

View File

@ -1,4 +1,4 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-
from .api import Galaxies, Clusters, EncodeGalaxies, UnableToRevertMachinetag
from .api import Galaxies, Clusters, EncodeGalaxies, EncodeClusters, UnableToRevertMachinetag

View File

@ -18,9 +18,14 @@ except ImportError:
class EncodeGalaxies(JSONEncoder):
def default(self, obj):
try:
return obj._json()
except AttributeError:
if isinstance(obj, Galaxy):
return obj.to_dict()
return JSONEncoder.default(self, obj)
class EncodeClusters(JSONEncoder):
def default(self, obj):
if isinstance(obj, (Cluster, ClusterValue, ClusterValueMeta)):
return obj.to_dict()
return JSONEncoder.default(self, obj)
@ -45,7 +50,7 @@ class Galaxy():
self.version = self.galaxy['version']
self.uuid = self.galaxy['uuid']
def _json(self):
def to_dict(self):
return {'type': self.type, 'name': self.name, 'description': self.description,
'version': self.version, 'uuid': self.uuid, 'icon': self.icon}
@ -104,7 +109,7 @@ class ClusterValueMeta():
# defined on the schema
self.additional_properties = m
def _json(self):
def to_dict(self):
to_return = {}
if self.type:
to_return['type'] = self.type
@ -160,12 +165,12 @@ class ClusterValue():
return None
return ClusterValueMeta(m)
def _json(self):
def to_dict(self):
to_return = {'value': self.value}
if self.description:
to_return['description'] = self.description
if self.meta:
to_return['meta'] = self.meta._json()
to_return['meta'] = self.meta
return to_return
@ -212,11 +217,11 @@ class Cluster(collections.Mapping):
def __iter__(self):
return iter(self.cluster_values)
def _json(self):
def to_dict(self):
to_return = {'name': self.name, 'type': self.type, 'source': self.source,
'authors': self.authors, 'description': self.description,
'uuid': self.uuid, 'version': self.version, 'values': []}
to_return['values'] = [v._json() for v in self.values()]
to_return['values'] = [v for v in self.values()]
return to_return

View File

@ -2,7 +2,7 @@
# -*- coding: utf-8 -*-
import unittest
from pymispgalaxies import Galaxies, Clusters, UnableToRevertMachinetag, EncodeGalaxies
from pymispgalaxies import Galaxies, Clusters, UnableToRevertMachinetag, EncodeGalaxies, EncodeClusters
from glob import glob
import os
import json
@ -22,7 +22,7 @@ class TestPyMISPGalaxies(unittest.TestCase):
galaxy = json.load(f)
galaxies_from_files[galaxy['name']] = galaxy
for name, g in self.galaxies.items():
out = g._json()
out = g.to_dict()
self.assertDictEqual(out, galaxies_from_files[g.name])
def test_dump_clusters(self):
@ -32,7 +32,7 @@ class TestPyMISPGalaxies(unittest.TestCase):
cluster = json.load(f)
clusters_from_files[cluster['name']] = cluster
for name, c in self.clusters.items():
out = c._json()
out = c.to_dict()
self.assertCountEqual(out, clusters_from_files[c.name])
def test_validate_schema_clusters(self):
@ -70,4 +70,4 @@ class TestPyMISPGalaxies(unittest.TestCase):
def test_json(self):
for c in self.clusters.values():
json.dumps(c, cls=EncodeGalaxies)
json.dumps(c, cls=EncodeClusters)