Improve code flow for Django 1.7/1.8 in patch_lookup args

polymorphic/query.py

@@ -182,55 +182,53 @@ class PolymorphicQuerySet(QuerySet):
     def _process_aggregate_args(self, args, kwargs):
         """for aggregate and annotate kwargs: allow ModelX___field syntax for kwargs, forbid it for args.
         Modifies kwargs if needed (these are Aggregate objects, we translate the lookup member variable)"""
-
-        def patch_lookup_lt_18(a):
-            a.lookup = translate_polymorphic_field_path(self.model, a.lookup)
-
-
-        def patch_lookup_gte_18(a):
-            # With Django > 1.8, the field on which the aggregate operates is
-            # stored inside a complex query expression.
-            if isinstance(a, Q):
-                translate_polymorphic_Q_object(self.model, a)
-            elif hasattr(a, 'get_source_expressions'):
-                for source_expression in a.get_source_expressions():
-                    if source_expression is not None:
-                        patch_lookup_gte_18(source_expression)
-            else:
-                a.name = translate_polymorphic_field_path(self.model, a.name)
-
         ___lookup_assert_msg = 'PolymorphicModel: annotate()/aggregate(): ___ model lookup supported for keyword arguments only'
-        def test___lookup_for_args_lt_18(a):
-            assert '___' not in a.lookup, ___lookup_assert_msg
 
-        def test___lookup_for_args_gte_18(a):
+        if django.VERSION < (1, 8):
-            """ *args might be complex expressions too in django 1.8 so
+            def patch_lookup(a):
-            the testing for a '___' is rather complex on this one """
+                a.lookup = translate_polymorphic_field_path(self.model, a.lookup)
-            if isinstance(a, Q):
-                def tree_node_test___lookup(my_model, node):
-                    " process all children of this Q node "
-                    for i in range(len(node.children)):
-                        child = node.children[i]
 
-                        if type(child) == tuple:
+            def test___lookup(a):
-                            # this Q object child is a tuple => a kwarg like Q( instance_of=ModelB )
+                assert '___' not in a.lookup, ___lookup_assert_msg
-                            assert '___' not in child[0], ___lookup_assert_msg
+        else:
-                        else:
+            def patch_lookup(a):
-                            # this Q object child is another Q object, recursively process this as well
+                # With Django > 1.8, the field on which the aggregate operates is
-                            tree_node_test___lookup(my_model, child)
+                # stored inside a complex query expression.
+                if isinstance(a, Q):
+                    translate_polymorphic_Q_object(self.model, a)
+                elif hasattr(a, 'get_source_expressions'):
+                    for source_expression in a.get_source_expressions():
+                        if source_expression is not None:
+                            patch_lookup(source_expression)
+                else:
+                    a.name = translate_polymorphic_field_path(self.model, a.name)
 
-                tree_node_test___lookup(self.model, a)
+            def test___lookup(a):
-            elif hasattr(a, 'get_source_expressions'):
+                """ *args might be complex expressions too in django 1.8 so
-                for source_expression in a.get_source_expressions():
+                the testing for a '___' is rather complex on this one """
-                    test___lookup_for_args_gte_18(source_expression)
+                if isinstance(a, Q):
-            else:
+                    def tree_node_test___lookup(my_model, node):
-                assert '___' not in a.name, ___lookup_assert_msg
+                        " process all children of this Q node "
+                        for i in range(len(node.children)):
+                            child = node.children[i]
+
+                            if type(child) == tuple:
+                                # this Q object child is a tuple => a kwarg like Q( instance_of=ModelB )
+                                assert '___' not in child[0], ___lookup_assert_msg
+                            else:
+                                # this Q object child is another Q object, recursively process this as well
+                                tree_node_test___lookup(my_model, child)
+
+                    tree_node_test___lookup(self.model, a)
+                elif hasattr(a, 'get_source_expressions'):
+                    for source_expression in a.get_source_expressions():
+                        test___lookup(source_expression)
+                else:
+                    assert '___' not in a.name, ___lookup_assert_msg
 
         for a in args:
-            test___lookup = test___lookup_for_args_lt_18 if django.VERSION < (1, 8) else test___lookup_for_args_gte_18
             test___lookup(a)
         for a in six.itervalues(kwargs):
-            patch_lookup = patch_lookup_lt_18 if django.VERSION < (1, 8) else patch_lookup_gte_18
             patch_lookup(a)
 
     def annotate(self, *args, **kwargs):