Merge pull request #53 from materialsproject/qe-result-map-fix

Fix QueryEngine result mapping for (nested) array fields

matgendb/query_engine.py

@@ -634,19 +634,23 @@ def _mapped_result(self, r):
             # Map aliased keys back to original key
             for k, v in self._prop_dict.items():
                 try:
-                    data = r[v[0]]
-                    for j in range(1, len(v)):
-                        if isinstance(data, list):
-                            data = [d[v[j]] for d in data]
-                        else:
-                            data = data[v[j]]
-                    result[k] = data
+                    result[k] = self._mapped_result_path(v[1:], data=r[v[0]])
                 except (IndexError, KeyError, ValueError):
                     result[k] = None
-
-
         return result
 
+    @staticmethod
+    def _mapped_result_path(path, data=None):
+        if not path:
+            return data
+        if isinstance(data, list):
+            return [QueryResults._mapped_result_path(path, d) for d in data]
+        else:
+            try:
+                return QueryResults._mapped_result_path(path[1:], data[path[0]])
+            except (IndexError, KeyError, ValueError):
+                return None
+
     def _result_generator(self):
         for r in self._results:
             yield self._mapped_result(r)

matgendb/tests/common.py

@@ -22,6 +22,11 @@
 
 _log = logging.getLogger('matgendb.tests')
 
+TEST_FILES_DIR = os.path.join(
+    os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
+    "test_files")
+
+
 def has_mongo():
     """Determine if MongoDB is up and usable
     """

matgendb/tests/test_queryresults.py

@@ -0,0 +1,62 @@
+import os
+import uuid
+import unittest
+
+import bson
+import pymongo
+
+from matgendb.query_engine import QueryEngine, QueryResults
+from matgendb.tests import common
+
+has_mongo = common.has_mongo()
+
+
+class QueryResultsTest(unittest.TestCase):
+    def setUp(self):
+        if has_mongo:
+            self.conn = pymongo.MongoClient()
+            self.db_name = 'test'
+            self.db = self.conn[self.db_name]
+            self.coll_name = 'tasks_{}'.format(uuid.uuid4())
+            self.coll = self.db[self.coll_name]
+            with open(os.path.join(common.TEST_FILES_DIR, 'db_test', 'GaLa.task.json')) as f:
+                doc = bson.json_util.loads(f.read())
+                self.coll.insert(doc)
+
+    def tearDown(self):
+        if has_mongo:
+            self.db.drop_collection(self.coll_name)
+
+    @unittest.skipUnless(has_mongo, 'requires MongoDB server')
+    def test_queryresult(self):
+        qe = QueryEngine(
+            connection=self.conn,
+            database=self.db_name,
+            collection=self.coll_name,
+        )
+        result = qe.query(
+            criteria={'task_id': 'mp-1002133'},
+            properties=[
+                'calcs_reversed.output.ionic_steps.e_0_energy',
+                'calcs_reversed.output.ionic_steps.electronic_steps.e_0_energy',
+            ],
+        )
+        self.assertTrue(isinstance(result, QueryResults))
+        print(list(qe.query(criteria={'task_id': 'mp-1002133'})))
+        self.assertEqual(len(result), 1)
+        doc = list(result)[0]
+        self.assertIn('calcs_reversed.output.ionic_steps.e_0_energy', doc)
+        v = doc['calcs_reversed.output.ionic_steps.e_0_energy']
+        self.assertIsInstance(v, list)
+        for elt in v:
+            self.assertIsInstance(elt, list)
+            for n in elt:
+                self.assertIsInstance(n, float)
+        self.assertIn('calcs_reversed.output.ionic_steps.electronic_steps.e_0_energy', doc)
+        v = doc['calcs_reversed.output.ionic_steps.electronic_steps.e_0_energy']
+        for elt in v:
+            self.assertIsInstance(elt, list)
+            for _elt in elt:
+                self.assertIsInstance(_elt, list)
+                for n in _elt:
+                    self.assertIsInstance(n, float)