Fix get_types for operation_types (refs #3179)

1 files changed, 22 insertions, 12 deletions

diff --git a/ishtar_common/models.py b/ishtar_common/models.py
index 04077b240..c536b64ae 100644
--- a/ishtar_common/models.py
+++ b/ishtar_common/models.py
@@ -329,29 +329,36 @@ class GeneralType(Cached, models.Model):
         return u""
 
     @classmethod
-    def get_types(cls, dct={}, instances=False, exclude=[], empty_first=True,
-                  default=None, initial=None):
-        types = cls.pre_get_types(dct, instances, exclude, empty_first,
-                                  default)
-        if not initial:
-            return types
+    def _get_initial_types(cls, initial, type_pks):
+        new_vals = []
         for value in initial:
             try:
                 pk = int(value)
             except ValueError:
                 continue
-            if pk in [idx for idx, lbl in types]:
+            if pk in type_pks:
                 continue
             try:
                 extra_type = cls.objects.get(pk=pk)
-                types.append((extra_type.pk, unicode(extra_type)))
+                new_vals.append((extra_type.pk, unicode(extra_type)))
             except cls.DoesNotExist:
                 continue
+        return new_vals
+
+    @classmethod
+    def get_types(cls, dct={}, instances=False, exclude=[], empty_first=True,
+                  default=None, initial=None):
+        types = cls._pre_get_types(dct, instances, exclude, empty_first,
+                                   default)
+        if not initial:
+            return types
+        new_vals = cls._get_initial_types(initial, [idx for idx, lbl in types])
+        types += new_vals
         return types
 
     @classmethod
-    def pre_get_types(cls, dct={}, instances=False, exclude=[],
-                      empty_first=True, default=None):
+    def _pre_get_types(cls, dct={}, instances=False, exclude=[],
+                       empty_first=True, default=None):
         # cache
         cache_key = None
         if not instances:
@@ -2637,7 +2644,7 @@ class OperationType(GeneralType):
 
     @classmethod
     def get_types(cls, dct={}, instances=False, exclude=[], empty_first=True,
-                  default=None):
+                  default=None, initial=[]):
         tuples = []
         dct['available'] = True
         if not instances and empty_first and not default:
@@ -2654,7 +2661,10 @@ class OperationType(GeneralType):
         if exclude:
             items = items.exclude(txt_idx__in=exclude)
         current_preventive, current_lst = None, None
-        for item in items.order_by(*cls._meta.ordering).all():
+        item_list = list(items.order_by(*cls._meta.ordering).all())
+        new_vals = cls._get_initial_types(initial, [i.pk for i in item_list])
+        item_list += new_vals
+        for item in item_list:
             if not current_lst or item.preventive != current_preventive:
                 if current_lst:
                     tuples.append(current_lst)