add color to Jupyter notebook module

pvviewer/__init__.py

@@ -7,6 +7,11 @@
 
 import uuid
 
+from viewer import Viewer
+import color
+import mol
+
+
 base_url = ('http://www.rcsb.org/pdb/download/downloadFile.do'
             '?fileFormat=pdb&structureId=')
 

pvviewer/color.py

@@ -0,0 +1,46 @@
+"""
+Contains the coloring operations
+"""
+
+from command import Command
+
+
+def _color_command(name, *args, **kwargs):
+    return Command('pv.color', name, args, kwargs)
+
+
+def _gradient(name):
+    return Command('pv.color', 'gradient', [name], {})
+
+
+def uniform(color='red'):
+    return _color_command('uniform', color)
+
+
+def by_element():
+    return _color_command('byElement')
+
+
+def by_ss():
+    return _color_command('bySS')
+
+
+def by_chain(gradient='rainbow'):
+    return _color_command('byChain', _gradient(gradient))
+
+
+def ss_succession(gradient='rainbow', coil_color='lightgrey'):
+    return _color_command('ssSuccession', _gradient(gradient), coil_color)
+
+
+def by_atom_prop(prop, gradient='rainbow', min_max=None):
+    return _color_command('byAtomProp', prop, _gradient(gradient), min_max)
+
+
+def by_residue_prop(prop, gradient='rainbow', min_max=None):
+    return _color_command('byResidueProp', prop, _gradient(gradient), min_max)
+
+
+
+
+

pvviewer/command.py

@@ -18,7 +18,7 @@ def encode(obj):
         # FIXME: Nan, Inf handling
         return str(obj)
     if isinstance(obj, dict):
-        contents = ['%s:%s' % (encode(k), encode(v)) for k,v in obj]
+        contents = ['%s:%s' % (encode(k), encode(v)) for k,v in obj.iteritems()]
         return '{%s}' % ', '.join(contents)
 
 
@@ -28,16 +28,22 @@ class Command:
     call that can be translated to JS.
     """
 
-    def __init__(self, receiver, command, args, kwargs):
+    def __init__(self, receiver, command, args, kwargs, terminate=False):
         self._receiver = receiver
         self._command = command
         self._args = args
         self._kwargs = kwargs
+        self._terminate = terminate
 
     def to_js(self):
         all_args = [', '.join(encode(arg) for arg in self._args)]
         if self._kwargs:
             all_args.append(encode(self._kwargs))
 
         args_string = ', '.join(all_args)
-        return '%s.%s(%s);' % (self._receiver, self._command, args_string)
+        t = self._terminate and ';' or ''
+        if not self._receiver:
+            call = self._command
+        else:
+            call = '%s.%s' % (self._receiver, self._command)
+        return '%s(%s)%s' % (call, args_string, t)

pvviewer/mol.py

@@ -38,6 +38,22 @@ def from_file(file_name):
         with open(file_name, 'r') as data:
             return Mol(data.read())
 
+
+def from_url(url):
+    return Mol.from_url(url)
+
+
+def from_file(file_name):
+    return Mol.from_file(file_name)
+
+
+def from_pdb_id(pdb_id):
+    return Mol.from_pdb_id(pdb_id)
+
+
 __all__ = (
     Mol,
+    from_file,
+    from_pdb_id,
+    from_url
 )

pvviewer/viewer.py

@@ -3,7 +3,7 @@
 notebooks.
 """
 
-import random
+import uuid
 from command import Command
 
 _VIEWER_SCAFFOLD_BEGIN = """
@@ -19,7 +19,7 @@
         viewer = pv.Viewer(document.getElementById('%(id)s'),
                            {quality : 'medium', width: 'auto',
                             height : 'auto', antialias : true,
-                            background : '#ddd',
+                            background : '#eee',
                             outline : true, style : '%(style)s' });
         viewer.fitParent();
 """
@@ -30,21 +30,14 @@
 """
 
 
-def _generate_id():
-    """
-    Generate a random id, suitable for use as a JS identifier.
-    """
-    lower_case_letters = 'abcdefghijklmnopqrstuvwxyz'
-    return '_%s' % ''.join(random.choice(lower_case_letters) for i in xrange(5))
-
-
 class Rendered:
     def __init__(self, data):
         self._data =data
 
     def _repr_html_(self):
         return self._data
 
+
 class Viewer:
     """
     This class implements the main protein viewer interface and lets you add
@@ -54,16 +47,14 @@ class Viewer:
     """
 
     def __init__(self):
-        # unique name for our parent div. Will come in handy later.
-        self._id = _generate_id()
         self._style = 'phong'
         self._width = 500
         self._height = 500
         self._commands = []
 
     def show(self):
         replacements = {
-            'id': self._id,
+            'id': str(uuid.uuid4()),
             'width': self._width,
             'height': self._height,
             'style': self._style