OpenStructure
color_tool.py

Shows how to write a tool in Python, register it and interact with the graphical objects in the scene.

1 """
2 This script illustrates how to implement a custom tool for interacting with the
3 3D scene.
4 """
5 from ost import gui
6 
7 
8 class ColorTool(gui.Tool):
9  def __init__(self):
10  gui.Tool.__init__(self, "Color Tool")
11  self.tool_options.AddOption(gui.ToolOptionFloat("r", "Red", 0.5, 0.0, 1.0))
12  self.tool_options.AddOption(gui.ToolOptionFloat("g", "Green", 0.5, 0.0, 1.0))
13  self.tool_options.AddOption(gui.ToolOptionFloat("b", "Blue", 0.5, 0.0, 1.0))
14  apply_to=gui.ToolOptionEnum("apply_to", "Apply To")
15  apply_to.Add("Atom", 0)
16  apply_to.Add("Residue", 1)
17  apply_to.Add("Chain", 2)
18  apply_to.SetIndex(0)
19  self.tool_options.AddOption(apply_to)
20 
21  def GetColor(self):
22  return gfx.Color(self.tool_options.GetOption('r').value,
23  self.tool_options.GetOption('g').value,
24  self.tool_options.GetOption('b').value)
25 
26  def CanOperateOn(self, obj):
27  return True
28  def GetIconPath(self):
29  return ''
30  def GetAtomsToColor(self, atom):
31  apply_to=self.tool_options.GetOption('apply_to').value
32  if apply_to==0:
33  return [atom]
34  if apply_to==1:
35  return atom.residue.atoms
36  if apply_to==2:
37  l=[]
38  for r in atom.residue.chain.residues:
39  for a in r.atoms:
40  l.append(a)
41  return l
42 
43  def Click(self, event):
44  obj, atom=gfx.PickAtom(scene, event.pos.x(), event.pos.y())
45  if atom.IsValid():
46  atoms=self.GetAtomsToColor(atom)
47  color=self.GetColor()
48  for atom in atoms:
49  obj.SetColorForAtom(color, atom.handle)
50 
51 color_tool=ColorTool()
52 
53 gui.ToolManager.Instance().AddTool(color_tool)
54 
55 e=io.LoadEntity(len(sys.argv)>1 and sys.argv[1] or '../entity/fragment.pdb')
56 g=gfx.Entity('e', e)
57 g.SetRenderMode(gfx.RenderMode.CUSTOM)
58 scene.Add(g)
59 scene.center=g.center