From 03b198598025d522aca1731d18029a8bcffa121b Mon Sep 17 00:00:00 2001
From: Konsta Vesterinen <konsta.vesterinen@gmail.com>
Date: Wed, 21 Aug 2013 15:20:31 +0300
Subject: [PATCH] Support for many-to-one batch fetching

---
 CHANGES.rst                                     |  6 ++++++
 setup.py                                        |  2 +-
 sqlalchemy_utils/__init__.py                    |  2 +-
 sqlalchemy_utils/functions/batch_fetch.py       | 12 +++++++++---
 .../test_many_to_one_relationships.py           | 17 +++++++++++++----
 5 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/CHANGES.rst b/CHANGES.rst
index e9ea715..10a3de6 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -4,6 +4,12 @@ Changelog
 Here you can see the full list of changes between each SQLAlchemy-Utils release.
 
 
+0.16.9 (2013-08-21)
+^^^^^^^^^^^^^^^^^^^
+
+- Support for many-to-one directed relationship properties batch fetching
+
+
 0.16.8 (2013-08-21)
 ^^^^^^^^^^^^^^^^^^^
 
diff --git a/setup.py b/setup.py
index 3ef8d17..b3f5bec 100644
--- a/setup.py
+++ b/setup.py
@@ -55,7 +55,7 @@ for name, requirements in extras_require.items():
 
 setup(
     name='SQLAlchemy-Utils',
-    version='0.16.8',
+    version='0.16.9',
     url='https://github.com/kvesteri/sqlalchemy-utils',
     license='BSD',
     author='Konsta Vesterinen',
diff --git a/sqlalchemy_utils/__init__.py b/sqlalchemy_utils/__init__.py
index 482d88d..2917856 100644
--- a/sqlalchemy_utils/__init__.py
+++ b/sqlalchemy_utils/__init__.py
@@ -37,7 +37,7 @@ from .types import (
 )
 
 
-__version__ = '0.16.8'
+__version__ = '0.16.9'
 
 
 __all__ = (
diff --git a/sqlalchemy_utils/functions/batch_fetch.py b/sqlalchemy_utils/functions/batch_fetch.py
index 7c8c1cf..1c8ccc8 100644
--- a/sqlalchemy_utils/functions/batch_fetch.py
+++ b/sqlalchemy_utils/functions/batch_fetch.py
@@ -156,7 +156,9 @@ class Fetcher(object):
         self.entities = coordinator.entities
         self.first = self.entities[0]
         self.session = object_session(self.first)
+        self.init_parent_dict()
 
+    def init_parent_dict(self):
         self.parent_dict = dict(
             (self.local_values(entity), [])
             for entity in self.entities
@@ -237,6 +239,12 @@ class ManyToManyFetcher(Fetcher):
 
 
 class ManyToOneFetcher(Fetcher):
+    def init_parent_dict(self):
+        self.parent_dict = dict(
+            (self.local_values(entity), None)
+            for entity in self.entities
+        )
+
     def fetch(self):
         column_name = list(self.prop.remote_side)[0].name
 
@@ -248,9 +256,7 @@ class ManyToOneFetcher(Fetcher):
         )
 
         for entity in self.related_entities:
-            self.parent_dict[getattr(entity, column_name)].append(
-                entity
-            )
+            self.parent_dict[getattr(entity, column_name)] = entity
 
 
 class OneToManyFetcher(Fetcher):
diff --git a/tests/batch_fetch/test_many_to_one_relationships.py b/tests/batch_fetch/test_many_to_one_relationships.py
index 36b0af4..4c774a0 100644
--- a/tests/batch_fetch/test_many_to_one_relationships.py
+++ b/tests/batch_fetch/test_many_to_one_relationships.py
@@ -28,17 +28,25 @@ class TestBatchFetchManyToOneRelationships(TestCase):
 
     def setup_method(self, method):
         TestCase.setup_method(self, method)
+        self.users = [
+            self.User(id=333, name=u'John'),
+            self.User(id=334, name=u'Matt')
+        ]
         articles = [
             self.Article(
                 id=1,
                 name=u'Article 1',
-                author=self.User(id=333, name=u'John')
+                author=self.users[0]
             ),
             self.Article(
                 id=2,
                 name=u'Article 2',
-                author=self.User(id=334, name=u'Matt')
+                author=self.users[1]
             ),
+            self.Article(
+                id=3,
+                name=u'Article 3'
+            )
         ]
         self.session.add_all(articles)
         self.session.commit()
@@ -50,6 +58,7 @@ class TestBatchFetchManyToOneRelationships(TestCase):
             'author'
         )
         query_count = self.connection.query_count
-        assert articles[0].author  # no lazy load should occur
-        assert articles[1].author  # no lazy load should occur
+        assert articles[0].author == self.users[0]  # no lazy load should occur
+        assert articles[1].author == self.users[1]  # no lazy load should occur
+        assert articles[2].author is None  # no lazy load should occur
         assert self.connection.query_count == query_count