From 6d351b21180d3e93b29fb9f376bf9fa9893be69e Mon Sep 17 00:00:00 2001
From: davebshow <davebshow@gmail.com>
Date: Sun, 3 Jul 2016 02:15:51 -0400
Subject: [PATCH] polishing

---
 goblin/api.py | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/goblin/api.py b/goblin/api.py
index 27df06a..1528da1 100644
--- a/goblin/api.py
+++ b/goblin/api.py
@@ -98,24 +98,20 @@ class Session:
         # Need some optional kwargs etc...
         while self._pending:
             elem = self._pending.popleft()
-            result = await self.save_element(elem)
-            if result:
-                self._current[result.id] = result
-
+            await self.save_element(elem)
+                
     async def save_element(self, element):
         if element.__type__ == 'vertex':
-            result = await self.save_vertex(element)
+            await self.save_vertex(element)
         elif element.__type__ == 'edge':
-            result = await self.save_edge(element)
-        else:
-            result = None
-        return result
+            await self.save_edge(element)
 
     async def save_vertex(self, element):
         result = await self._save_element(element,
                                           self._create_vertex,
                                           self._update_vertex,
                                           mapper.map_vertex_to_ogm)
+        self._current[result.id] = result
         return result
 
     async def save_edge(self, element):
@@ -125,6 +121,7 @@ class Session:
                                           self._create_edge,
                                           self._update_edge,
                                           mapper.map_edge_to_ogm)
+        self._current[result.id] = result
         return result
 
     async def _save_element(self,
-- 
GitLab