From afa4e984eb3a2e5b3f67262dc293bd9a97841de1 Mon Sep 17 00:00:00 2001
From: Konsta Vesterinen <konsta.vesterinen@gmail.com>
Date: Sat, 4 Jan 2014 02:43:38 +0200
Subject: [PATCH] Remove CompositePath

---
 sqlalchemy_utils/batch.py                   | 43 +++++++++------------
 tests/batch_fetch/test_compound_fetching.py |  5 +--
 2 files changed, 21 insertions(+), 27 deletions(-)

diff --git a/sqlalchemy_utils/batch.py b/sqlalchemy_utils/batch.py
index bbbca33..f4f5e40 100644
--- a/sqlalchemy_utils/batch.py
+++ b/sqlalchemy_utils/batch.py
@@ -27,6 +27,10 @@ class PathException(Exception):
     pass
 
 
+class DataException(Exception):
+    pass
+
+
 class with_backrefs(object):
     """
     Marks given attribute path so that whenever its fetched with batch_fetch
@@ -37,9 +41,6 @@ class with_backrefs(object):
         self.path = path
 
 
-
-
-
 class Path(object):
     """
     A class that represents an attribute path.
@@ -79,7 +80,8 @@ class Path(object):
                     related_entities.extend(getattr(entity, attrs[0]))
 
                 if not related_entities:
-                    return
+                    raise DataException('No related entities.')
+
                 subpath = '.'.join(attrs[1:])
                 return Path.parse(related_entities, subpath, populate_backrefs)
             else:
@@ -107,11 +109,6 @@ class Path(object):
                     return OneToManyFetcher
 
 
-class CompositePath(object):
-    def __init__(self, *paths):
-        self.paths = paths
-
-
 def batch_fetch(entities, *attr_paths):
     """
     Batch fetch given relationship attribute for collection of entities.
@@ -174,10 +171,16 @@ def batch_fetch(entities, *attr_paths):
 
     if entities:
         for path in attr_paths:
-            fetcher = fetcher_factory(entities, path)
-            if fetcher:
+            try:
+                fetcher = fetcher_factory(entities, path)
                 fetcher.fetch()
                 fetcher.populate()
+            except DataException:
+                pass
+
+
+def get_fetcher(entities, path, populate_backrefs):
+    return Path.parse(entities, path, populate_backrefs).fetcher
 
 
 def fetcher_factory(entities, path):
@@ -186,20 +189,12 @@ def fetcher_factory(entities, path):
         path = path.path
         populate_backrefs = True
 
-    if isinstance(path, CompositePath):
-        fetchers = []
-        for path in path.paths:
-            path = Path.parse(entities, path, populate_backrefs)
-            if path:
-                fetchers.append(
-                    path.fetcher
-                )
-
-        return CompositeFetcher(*fetchers)
+    if isinstance(path, tuple):
+        return CompositeFetcher(
+            *(get_fetcher(entities, p, populate_backrefs) for p in path)
+        )
     else:
-        path = Path.parse(entities, path, populate_backrefs)
-        if path:
-            return path.fetcher
+        return get_fetcher(entities, path, populate_backrefs)
 
 
 class CompositeFetcher(object):
diff --git a/tests/batch_fetch/test_compound_fetching.py b/tests/batch_fetch/test_compound_fetching.py
index 7f74a3c..db24c76 100644
--- a/tests/batch_fetch/test_compound_fetching.py
+++ b/tests/batch_fetch/test_compound_fetching.py
@@ -1,6 +1,5 @@
 import sqlalchemy as sa
 from sqlalchemy_utils import batch_fetch
-from sqlalchemy_utils.batch import CompositePath
 from tests import TestCase
 
 
@@ -93,7 +92,7 @@ class TestCompoundOneToManyBatchFetching(TestCase):
         batch_fetch(
             buildings,
             'business_premises',
-            CompositePath(
+            (
                 'equipment',
                 'business_premises.equipment'
             )
@@ -198,7 +197,7 @@ class TestCompoundManyToOneBatchFetching(TestCase):
         batch_fetch(
             buildings,
             'business_premises',
-            CompositePath(
+            (
                 'equipment',
                 'business_premises.equipment'
             )