Fixed search when haystack spans multiple objects and the same column is selected in all objects

Author: Evgenii Osipov

findseq.py

@@ -70,24 +70,6 @@
 
 
 def findseq(needle, haystack, selName=None, het=0, firstOnly=0):
-    """USAGE:
-        findseq needle, haystack[, selName[, het[, firstOnly]]]
-    """
-    # set the name of the selection to return.
-    if selName == None:
-        rSelName = "foundSeq" + str(random.randint(0, 32000))
-        selName = rSelName
-    elif selName == "sele":
-        rSelName = "sele"
-    else:
-        rSelName = selName
-
-    # input checking
-    if not checkParams(needle, haystack, selName, het, firstOnly):
-        print("There was an error with a parameter.  Please see")
-        print("the above error message for how to fix it.")
-        return None
-
     one_letter = {
         '00C': 'C', '01W': 'X', '0A0': 'D', '0A1': 'Y', '0A2': 'K',
         '0A8': 'C', '0AA': 'V', '0AB': 'V', '0AC': 'G', '0AD': 'G',
@@ -338,51 +320,62 @@ def findseq(needle, haystack, selName=None, het=0, firstOnly=0):
         'YG ': 'G', 'YOF': 'Y', 'YRR': 'N', 'YYG': 'G', 'Z  ': 'C',
         'ZAD': 'A', 'ZAL': 'A', 'ZBC': 'C', 'ZCY': 'C', 'ZDU': 'U',
         'ZFB': 'X', 'ZGU': 'G', 'ZHP': 'N', 'ZTH': 'T', 'ZZJ': 'A'}
-
-    # remove hetero atoms (waters/ligands/etc) from consideration?
-    if het:
-        cmd.select("__h", "br. " + haystack)
+    # set the name of the selection to return.
+    if selName == None:
+        rSelName = "foundSeq" + str(random.randint(0, 32000))
+        selName = rSelName
+    elif selName == "sele":
+        rSelName = "sele"
     else:
-        cmd.select("__h", "br. " + haystack + " and not het")
+        rSelName = selName
+    # make an empty selection to which we add residues
+    cmd.select(rSelName, 'None')
+    for obj in cmd.get_object_list(haystack):
+        # input checking
+        if not checkParams(needle, haystack, selName, het, firstOnly):
+            print("There was an error with a parameter.  Please see")
+            print("the above error message for how to fix it.")
+            return None
+
 
-    # get the AAs in the haystack
-    aaDict = {'aaList': []}
-    cmd.iterate("(name ca) and __h", "aaList.append((resi,resn,chain))", space=aaDict)
+        # remove hetero atoms (waters/ligands/etc) from consideration?
+        if het:
+            cmd.select("__h", f"br. {obj} and {haystack}")
+        else:
+            cmd.select("__h", f"br. {obj} and {haystack} and not het")
 
-    IDs = [x[0] for x in aaDict['aaList']]
-    AAs = ''.join([one_letter[x[1]] for x in aaDict['aaList']])
-    chains = [x[2] for x in aaDict['aaList']]
+        # get the AAs in the haystack
+        aaDict = {'aaList': []}
+        cmd.iterate("(name ca) and __h", "aaList.append((resi,resn,chain))", space=aaDict)
 
-    reNeedle = re.compile(needle.upper())
-    it = reNeedle.finditer(AAs)
+        IDs = [x[0] for x in aaDict['aaList']]
+        AAs = ''.join([one_letter[x[1]] for x in aaDict['aaList']])
+        chains = [x[2] for x in aaDict['aaList']]
 
-    # make an empty selection to which we add residues
-    cmd.select(rSelName, 'None')
+        reNeedle = re.compile(needle.upper())
+        it = reNeedle.finditer(AAs)
 
-    for i in it:
-        (start, stop) = i.span()
-        # we found some residues, which chains are they from?
-        i_chains = chains[start:stop]
-        # are all residues from one chain?
-        if len(set(i_chains)) != 1:
-            # now they are not, this match is not really a match, skip it
-            continue
-        chain = i_chains[0]
-        # Only apply chains to selection algebra if there are defined chains.
-        if chain:
-            cmd.select(rSelName, rSelName + " or (__h and i. " + '+'.join(IDs[ii] for ii in range(start, stop)) + " and c. " + chain + " )")
-        else:
-            cmd.select(rSelName, rSelName + " or (__h and i. " + '+'.join(IDs[ii] for ii in range(start, stop)) + ")")
-        if int(firstOnly):
-            break
+
+        for i in it:
+            (start, stop) = i.span()
+            # we found some residues, which chains are they from?
+            i_chains = chains[start:stop]
+            # are all residues from one chain?
+            if len(set(i_chains)) != 1:
+                # now they are not, this match is not really a match, skip it
+                continue
+            chain = i_chains[0]
+            # Only apply chains to selection algebra if there are defined chains.
+            if chain:
+                cmd.select(rSelName, rSelName + " or (__h and i. " + '+'.join(IDs[ii] for ii in range(start, stop)) + " and c. " + chain + " )")
+            else:
+                cmd.select(rSelName, rSelName + " or (__h and i. " + '+'.join(IDs[ii] for ii in range(start, stop)) + ")")
+            if int(firstOnly):
+                break
     cmd.delete("__h")
     return rSelName
-
 cmd.extend("findseq", findseq)
-cmd.findseq = findseq
-cmd.auto_arg[1]['findseq'] = [ cmd.object_sc, 'object', '']
-cmd.auto_arg[2]['findseq'] = [lambda: cmd.Shortcut(['het=1','firstOnly=1']), 'params', '']
-cmd.auto_arg[3]['findseq'] = [lambda: cmd.Shortcut(['het=1','firstOnly=1']), 'params', '']
+
 
 def checkParams(needle, haystack, selName, het, firstOnly):
     """