Re-add structures

atomium/__init__.py

@@ -1,4 +1,5 @@
 from .utilities import open, fetch
+from .structures import Atom, Residue, Ligand, Chain, Model
 
 __author__ = "Sam Ireland"
 __version__ = "0.12.0"

atomium/base.py

@@ -0,0 +1,82 @@
+"""Decorators and metaclasses used by atomium structures."""
+
+import re
+
+def filter_objects(objects, key, value):
+    """Takes a dictionary of objects, and filters them on object properties.
+
+    :param dict objects: the dictionary of objects - the keys are unimportant.
+    :param str key: the attribute to search. This can be an attribute of the\
+    object, or attr__regex, or attr__gt etc.
+    :param value: the value that the attribute must have.
+    :rtype: ``dict``"""
+
+    if "__" in key:
+        attr, comp = key.split("__")
+        if comp == "regex":
+            objects = {i: o for i, o in objects.items()
+             if re.match(value, getattr(o, attr))}
+        else:
+            comp = "__{}__".format(comp)
+            objects = {i: o for i, o in objects.items()
+             if getattr(getattr(o, attr), comp)(value)}
+    else:
+        objects = {i: o for i, o in objects.items()
+         if getattr(o, key).__eq__(value)}
+    return objects
+
+
+def query(func, tuple_=False):
+    """A decorator which can be applied to any function which returns a ``dict``
+    and which takes no other paramters other than ``self``. It will query the
+    returned objects by any keyword argument, or use a positional argument to
+    search by ID.
+
+    :param func: the function to modify.
+    :param bool tuple_: if ``True``, objects will be returned in a tuple not a\
+    set.
+    :rtype: ``function``"""
+
+    def structures(self, *args, **kwargs):
+        objects = func(self)
+        if len(args) == 1:
+            return {objects[args[0]]} if args[0] in objects else set()
+        for k, v in kwargs.items():
+            objects = filter_objects(objects, k, v)
+        t = tuple if tuple_ else set
+        return t(objects.values())
+    return structures
+
+
+def getone(func):
+    """A decorator which can be applied to any function which returns an
+    iterable. It produces a function which just gets the first item in that
+    iterable.
+
+    :param func: the function to modify.
+    :rtype: ``function``"""
+
+    def structure(self, *args, **kwargs):
+        for obj in func(self, *args, **kwargs): return obj
+    return structure
+
+
+
+class StructureClass(type):
+    """A metaclass which can be applied to structure class. It will override
+    the instantation behaviour so that all methods that belong to a preset
+    list ('atoms', 'chains' etc.) will have the :py:func:`.query` decorator
+    applied and a copy with the :py:func:`.getone` decorator applied."""
+
+    METHODS = ["chains", "residues", "ligands", "waters", "molecules", "atoms"]
+
+    def __new__(self, *args, **kwargs):
+        cls = type.__new__(self, *args, **kwargs)
+        for attribute in dir(cls):
+            if attribute in cls.METHODS:
+                setattr(cls, attribute, query(
+                 getattr(cls, attribute),
+                 tuple_=(attribute == "residues" and cls.__name__ == "Chain")
+                ))
+                setattr(cls, attribute[:-1], getone(getattr(cls, attribute)))
+        return cls

atomium/data.py

@@ -1,5 +1,7 @@
 """Contains logic for turning data dictionaies into a parsed Python objects."""
 
+from .structures import *
+
 class File:
     """When a file is parsed, the result is a ``File``. It contains the
     structure of interest, as well as meta information.
@@ -142,6 +144,24 @@ def assemblies(self):
         return self._assemblies
 
 
+    @property
+    def models(self):
+        """The structure's models.
+
+        :rtype: ``list``"""
+
+        return self._models
+
+
+    @property
+    def model(self):
+        """The structure's first model (and only model if it has only one).
+
+        :rtype: ``Model``"""
+
+        return self._models[0]
+
+
 
 def data_dict_to_file(data_dict, filetype):
     """Turns an atomium data dictionary into a :py:class:`.File`.
@@ -155,9 +175,99 @@ def data_dict_to_file(data_dict, filetype):
         if key != "models":
             for subkey, value in data_dict[key].items():
                 setattr(f, "_" + subkey, value)
+    f._models = [model_dict_to_model(m) for m in data_dict["models"]]
     return f
 
 
+def model_dict_to_model(model_dict):
+    """Takes a model dictionary and turns it into a fully processed
+    :py:class:`.Model` object.
+
+    :param dict model_dict: the model dictionary.
+    :rtype: ``Model``"""
+
+    chains = create_chains(model_dict)
+    ligands = create_ligands(model_dict, chains)
+    waters = create_ligands(model_dict, chains, water=True)
+    model = Model(*(chains + ligands + waters))
+    return model
+
+
+def create_chains(model_dict):
+    """Creates a list of :py:class:`.Chain` objects from a model dictionary.
+
+    :param dict model_dict: the model dictionary.
+    :rtype: ``list``"""
+
+    chains = []
+    for chain_id, chain in model_dict["polymer"].items():
+        res = [create_het(r, i) for i, r in chain["residues"].items()]
+        for res1, res2 in zip(res[:-1], res[1:]):
+            res1._next, res2._previous = res2, res1
+        chains.append(Chain(*res, id=chain_id,
+         internal_id=chain["internal_id"], sequence=chain["sequence"]))
+    return chains
+
+
+def create_ligands(model_dict, chains, water=False):
+    """Creates a list of :py:class:`.Ligand` objects from a model dictionary.
+
+    :param dict model_dict: the model dictionary.
+    :param list chains: a list of :py:class:`.Chain` objects to assign by ID.
+    :param bool water: if `True``, water ligands will be made.
+    :rtype: ``list``"""
+
+    ligs = []
+    for lig_id, lig in model_dict["water" if water else "non-polymer"].items():
+        chain = None
+        for c in chains:
+            if c._id == lig["polymer"]:
+                chain = c
+                break
+        ligs.append(create_het(lig, lig_id, ligand=True, chain=chain, water=water))
+    return ligs
+
+
+def create_het(d, id, ligand=False, chain=None, water=False):
+    """Creates a :py:class:`.Residue` or :py:class:`.Ligand` from some
+    atom-containing dictionary.
+
+    If there is multiple occupancy, only one position will be used.
+
+    :param dict d: the dictionary to parse.
+    :param str id: the ID of the structure to make.
+    :param bool ligand: if ``True`` a ligand will be made, not a residue.
+    :param Chain chain: the :py:class:`.Chain` to assign if a ligand.
+    :param bool water: if ``True``, the ligand will be a water ligand.
+    :rtype: ``Residue`` or ``Ligand``"""
+
+    alt_loc = None
+    if any([atom["occupancy"] < 1 for atom in d["atoms"].values()]):
+        if any([atom["alt_loc"] for atom in d["atoms"].values()]):
+            alt_loc = sorted([atom["alt_loc"] for atom in d["atoms"].values()
+             if atom["alt_loc"]])[0]
+    atoms = [atom_dict_to_atom(a, i) for i, a in d["atoms"].items()
+     if a["occupancy"] == 1 or a["alt_loc"] is None or a["alt_loc"] == alt_loc]
+    if ligand:
+        return Ligand(*atoms, id=id, name=d["name"], chain=chain,
+         internal_id=d["internal_id"], water=water)
+    else:
+        return Residue(*atoms, id=id, name=d["name"])
+
+
+def atom_dict_to_atom(d, atom_id):
+    """Creates an :py:class:`.Atom` from an atom dictionary.
+
+    :param dict d: the atom dictionary.
+    :param int id: the atom's ID.
+    :rtype: ``Atom``"""
+
+    return Atom(
+     d["element"], d["x"], d["y"], d["z"], atom_id,
+     d["name"], d["charge"], d["bvalue"], d["anisotropy"]
+    )
+
+
 
 PERIODIC_TABLE = {
  "H": 1.0079, "HE": 4.0026, "LI": 6.941, "BE": 9.0122, "B": 10.811,