select with duplicates from different sources

Author: enriquepablo

src/pyff/builtins.py

@@ -832,7 +832,10 @@ def select(req: Plumbing.Request, *opts):
         if opts[0] == 'as' and len(opts) == 2:
             name = opts[1]
 
-    entities = resolve_entities(args, lookup_fn=req.md.store.select, dedup=dedup)
+    if dedup:
+        entities = resolve_entities(args, lookup_fn=req.md.store.select, dedup=dedup)
+    else:
+        entities = resolve_entities(args, lookup_fn=req.md.store.select_with_dups, dedup=dedup)
 
     if req.state.get('match', None):  # TODO - allow this to be passed in via normal arguments
 

src/pyff/store.py

@@ -4,6 +4,7 @@
 import re
 import shutil
 import time
+from collections import defaultdict
 from datetime import datetime, timedelta
 from io import BytesIO
 from threading import ThreadError
@@ -807,6 +808,7 @@ def __init__(self, *args, **kwargs):
         self.md = dict()
         self.index = dict()
         self.entities = dict()
+        self.md_entities = defaultdict(dict)
 
         for hn in DINDEX:
             self.index.setdefault(hn, {})
@@ -885,13 +887,15 @@ def update(self, t, tid=None, etag=None, lazy=True):
             self.entities[relt.get('entityID')] = relt  # TODO: merge?
             if tid is not None:
                 self.md[tid] = [relt.get('entityID')]
+                self.md_entities[tid][relt.get('entityID')] = relt
         elif relt.tag == "{%s}EntitiesDescriptor" % NS['md']:
             if tid is None:
                 tid = relt.get('Name')
             lst = []
             for e in iter_entities(t):
-                self.update(e)
+                self.update(e, tid)
                 lst.append(e.get('entityID'))
+                self.md_entities[tid][e.get('entityID')] = e
             self.md[tid] = lst
 
     def lookup(self, key):
@@ -946,3 +950,9 @@ def _lookup(self, key):
             return lst
 
         return []
+
+    def select_with_dups(self, member):
+        if member in self.md_entities:
+            return self.md_entities[member].values()
+
+        return []