CollapseSel

Overview

CollapseSel is a small utility function to compress strings for selections. So, if you have a selection with residues 1+2+3+4+20+21+22+100-120 this will return, 1-4+20-22+100-120.

CollapseIDs is a small utility function to compress strings for an array of IDs. This does NOT have the logic for detecting duplicates, CollapseSel does.

The Code

import pymol
from pymol import stored

def collapseIDs(ids):
	"""
	Helper function to make a smart list of IDs: eg turns 1+2+3+4+5 into 1-5.
	"""
	rVal = []
	if len(ids)==0:
		return ""

	scanning=False
	anchor = 0
	start = 0
	# 1-5 7-10 12 21-23
	for cur in range(0,len(ids)-1):
		if ids[cur]+1 == ids[cur+1]:
			if scanning:
				scanning=True
				continue
			else:
				scanning=True
				start = cur
		else:
			if scanning:
				scanning=False
				rVal.append(str(ids[start]) + "-" + str(ids[cur]))
				start = cur
			else:
				scanning=False
				rVal.append(str(ids[cur]))
	if scanning:
		rVal.append( str(ids[start]) + "-" + str(ids[cur+1]))
	else:
		rVal.append(str(ids[-1]))
	return rVal



def collapseSel(sel=None,lType="resi"):
	"""
	collapseSel -- given a valid PyMOL selection and list type, return a collapsed
		list of numbers corresponding to the lType.  For example, to compactly
		print the residue identifiers for all the waters, try:
			select theWaters, resn HOH
			print collapseSel( "theWaters" )
	
		This will convert: 1+2+3+4+10+11+12 into 1-4+10-12

	PARAMS
		sel
			The selection to collapse
		lType
			The identifier type: 'resi', 'ID', 'id', or any numerical property.

	RETURNS
		a string of collapsed IDs.
	"""
	if sel==None:
		return ""

	stored.s = set()
	cmd.iterate( sel, "stored.s.add(int(%s))" % lType)
	return "+".join(collapseIDs(list(stored.s)))

cmd.extend("collapseSel", collapseSel)
cmd.extend("collapsIDs", collapseIDs)