import copy
import random
from schrodinger.Qt import QtCore
ORIG_CONNECTION_COLORS = [[238, 162, 173], [191, 62, 255], [78, 238, 148],
[192, 192, 192], [72, 209, 204], [255, 193, 37],
[197, 193, 170], [202, 225, 255], [113, 113, 198],
[238, 238, 0]]
connection_colors = copy.deepcopy(ORIG_CONNECTION_COLORS)
[docs]def get_connection_color(i):
"""
Return the appropriate color for the i-th connection
:param i: The connection number to get colors for
:type i: int
:return: The appropriate color, formatted as a list of three integers
between 0 and 255
:rtype: list
"""
extend_connection_colors(i)
return copy.copy(connection_colors[i])
# Return a copy so the calling scope can't modify connection_colors
[docs]def get_color_list(i):
"""
Return a list of colors from 0 to i
:param i: The connection number to get colors up to
:type i: int
:return: A list of appropriate colors, where each color is formatted as a
list of three integers between 0 and 255
:rtype: list
"""
extend_connection_colors(i)
return copy.deepcopy(connection_colors[:i])
# Return a copy so the calling scope can't modify connection_colors
[docs]def get_flat_color_list(i):
"""
Return a flat list of colors from 0 to i
:param i: The connection number to get colors up to
:type i: int
:return: A flat list of appropriate colors of length 3*i
:rtype: list
"""
extend_connection_colors(i)
flat_list = []
list(map(flat_list.extend, connection_colors[:i]))
return flat_list
# We don't need to return a copy here since nothing in flat_list is linked
# to connection_colors
[docs]def extend_connection_colors(idx):
"""
Make sure that connection_colors is at least as long as the requested length
:param i: The desired length of connection_colors
:type i: int
"""
num_to_extend = idx - len(connection_colors) + 1
if num_to_extend <= 0:
return
new_colors = [gen_random_color() for i in range(num_to_extend)]
connection_colors.extend(new_colors)
[docs]def gen_random_color():
"""
Generate a new random color
:return: A random color, formatted as a list of three integers
between 0 and 255
:rtype: list
"""
return [random.randint(0, 255) for i in range(3)]
[docs]def reset_colors():
"""
Reset the connection_colors list back to its original state. Intended for
use in unit testing
"""
global connection_colors
connection_colors = copy.deepcopy(ORIG_CONNECTION_COLORS)
[docs]class ConnectionStorage(QtCore.QObject):
"""
This class is used to store the data of the current connection table.
"""
[docs] def __init__(self):
super(ConnectionStorage, self).__init__()
self._data = {}
self._changed = False
self.ignorebp = False
def __iter__(self):
tmpList = []
for bp in list(self._data):
tmpList.append((self.getRow(bp), bp))
tmpList.sort()
for row, bp in tmpList:
yield bp
[docs] def reset(self):
self._data = {}
self._changed = False
self.ignorebp = False
[docs] def collectionChanged(self):
self._changed = True
[docs] def hasChanged(self):
return self._changed
[docs] def changesSaved(self):
self._changed = False
[docs] def addPair(self,
bp,
name="",
csize=0,
minlink=0,
maxlink=0,
cfile=0,
bondToH=False,
selected=False):
"""
Adds a bond pair to the list of bond pairs, and also creates a
more or less empty data line for that data. Data in array:
[Attachment name, collection size, min linkers, max linkers,
collection file, table row]
"""
row = len(self._data)
self._data[bp] = [0, 0, 0, 0, 0, 0]
self.setData(bp, name, csize, minlink, maxlink, cfile, row, bondToH,
selected)
self._changed = True
return bp
[docs] def bondPair(self, a1, a2):
if (a1, a2) in self:
return (a1, a2)
if (a2, a1) in self:
return (a2, a1)
return None
[docs] def switchPair(self, bp, newbp=False):
if not newbp:
newbp = (bp[1], bp[0])
self._data[newbp] = self._data[bp]
if bp != newbp:
del self._data[bp]
self._changed = True
[docs] def changePair(self, oldbp, newbp):
changeddata = self._data[newbp]
self._data[newbp] = self._data[oldbp]
self._data[oldbp] = changeddata
self._changed = True
[docs] def deletePair(self, bp):
row = self.getRow(bp)
del self._data[bp]
for bp2 in self._data:
if self.getRow(bp2) > row:
self.setRow(bp2, self.getRow(bp2) - 1)
self._changed = True
[docs] def setData(self, bp, name, csize, minlink, maxlink, cfile, row, bondToH,
selected):
self._data[bp] = [
name, csize, minlink, maxlink, cfile, row, bondToH, selected
]
self._changed = True
[docs] def setConnectionName(self, bp, cname):
self._data[bp][0] = cname
self._changed = True
[docs] def setConnectionSize(self, bp, csize):
self._data[bp][1] = csize
self._changed = True
[docs] def setConnectionMinLink(self, bp, minlink):
self._data[bp][2] = minlink
self._changed = True
[docs] def setConnectionMaxLink(self, bp, maxlink):
self._data[bp][3] = maxlink
self._changed = True
[docs] def setConnectionFile(self, bp, cfile):
self._data[bp][4] = cfile
self._changed = True
[docs] def setRow(self, bp, row):
self._data[bp][5] = row
self._changed = True
[docs] def setSelected(self, bp, selected):
self._data[bp][7] = selected
[docs] def getConnectionName(self, bp):
return self._data[bp][0]
[docs] def getConnectionSize(self, bp):
return self._data[bp][1]
[docs] def getConnectionMinLink(self, bp):
return self._data[bp][2]
[docs] def getConnectionMaxLink(self, bp):
return self._data[bp][3]
[docs] def getConnectionFile(self, bp):
return self._data[bp][4]
[docs] def getRow(self, bp):
return self._data[bp][5]
[docs] def bondToH(self, bp):
return self._data[bp][6]
[docs] def getSelected(self, bp):
return self._data[bp][7]
[docs] def clearSelected(self):
for bp in self._data:
self.setSelected(bp, False)
[docs] def size(self):
return len(self._data)
[docs] def getPairFromRow(self, num):
for bp in self._data:
if self.getRow(bp) == num:
return bp
return False
[docs] def getBondColorList(self):
colorList = []
multiple_cores = self.containsMultipleCores()
for bp in self:
if self.bondToH(bp) or bp == self.ignorebp:
continue
i = self.getRow(bp)
if multiple_cores:
i -= 1
colorList += self.getConnectionColor(i)
return colorList
[docs] def containsMultipleCores(self):
bp = self.getPairFromRow(0)
name = self.getConnectionName(bp)
return (name == "Core")
[docs] def getConnectionColor(self, i):
"""
Return the appropriate connection color for the i-th connection
:param i: The connection number to get the color for
:type i: int
:return: The appropriate color, formatted as a list of three integers
between 0 and 255
:rtype: list
"""
return get_connection_color(i)
[docs] def setIgnorePair(self, bp):
self.ignorebp = bp
[docs] def removeIgnorePair(self):
self.ignorebp = False
[docs] def getBondList(self):
retList = []
for bp in self:
if self.bondToH(bp) or bp == self.ignorebp:
continue
retList.append(bp[1])
retList.append(bp[0])
retList.append(int(self.getSelected(bp)))
return retList
[docs] def getBondListWithIgnored(self):
retList = []
for bp in self:
retList.append(bp[1])
retList.append(bp[0])
retList.append(int(self.getSelected(bp)))
return retList
[docs] def getAtomList(self):
atomList = []
for bp in self:
if self.bondToH(bp) and bp != self.ignorebp:
atomList.append(bp[0])
return atomList
[docs] def getCompleteBondList(self):
atomList = []
for bp in self:
atomList.append(bp[0])
atomList.append(bp[1])
return atomList
[docs] def getAtomColorList(self):
colorList = []
for bp in self:
if not self.bondToH(bp) or bp == self.ignorebp:
continue
i = self.getRow(bp)
colorList += self.getConnectionColor(i)
return colorList