From 12f5302eb3c426c73ba3d48b85df75626955d293 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= <raphael@vinot.info>
Date: Wed, 1 Nov 2017 11:20:09 -0700
Subject: [PATCH] Properly use the JSONEncoders

---
 pymispgalaxies/__init__.py |  2 +-
 pymispgalaxies/api.py      | 25 +++++++++++++++----------
 tests/tests.py             |  8 ++++----
 3 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/pymispgalaxies/__init__.py b/pymispgalaxies/__init__.py
index 139bff8..d6ca4a5 100644
--- a/pymispgalaxies/__init__.py
+++ b/pymispgalaxies/__init__.py
@@ -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
diff --git a/pymispgalaxies/api.py b/pymispgalaxies/api.py
index 40009d0..0784421 100644
--- a/pymispgalaxies/api.py
+++ b/pymispgalaxies/api.py
@@ -18,10 +18,15 @@ except ImportError:
 
 class EncodeGalaxies(JSONEncoder):
     def default(self, obj):
-        try:
-            return obj._json()
-        except AttributeError:
-            return JSONEncoder.default(self, obj)
+        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)
 
 
 class PyMISPGalaxiesError(Exception):
@@ -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
 
 
diff --git a/tests/tests.py b/tests/tests.py
index bd88620..d08c393 100644
--- a/tests/tests.py
+++ b/tests/tests.py
@@ -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)