Source code for atomium.base

"""Decorators and metaclasses used by atomium structures."""

import re

[docs]def get_object_from_filter(obj, components): """Gets the object whose attributes are actually being queried, which may be a different object if there is a chain. :param obj: the intial object. :param list components: the components of the original key. :returns: the relevant object""" components = components[:] while len(components) > 2: obj = getattr(obj, components.pop(0)) if len(components) == 2: if components[-1] != "regex": if not hasattr(obj, f"__{components[-1]}__"): obj = getattr(obj, components[0]) return obj
[docs]def get_object_attribute_from_filter(obj, components): """Gets the object's value of some attribute based on a list of key components. :param obj: the object with attributes. :param list components: the components of the original key. :returns: the value""" try: return getattr( obj, components[-1] if hasattr(obj, components[-1]) else components[-2] ) except: return None
[docs]def attribute_matches_value(attribute, value, components): """Checks if an attribute value matches a given value. The components given will determine whether an exact match is sought, or whether a more complex criterion is used. :param attribute: the value of an object's attribute. :param value: the value to match against. :param list components: the components of the original key. :rtype: ``bool``""" if components[-1] == "regex": return re.match(value, attribute) possible_magic = f"__{components[-1]}__" if hasattr(attribute, possible_magic): return getattr(attribute, possible_magic)(value) return getattr(attribute, "__eq__")(value)
[docs]def filter_objects(objects, key, value): """Takes a :py:class:`.StructureSet` of objects, and filters them on object properties. They key can be an attribute of the object, or a complex double-underscore separated chain of attributes. :param StructreSet 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``""" components = key.split("__") matching_objects = [] for structure in objects.structures: obj = get_object_from_filter(structure, components) attr = get_object_attribute_from_filter(obj, components) if attribute_matches_value(attr, value, components): matching_objects.append(structure) return StructureSet(*matching_objects)
[docs]def query(func, tuple_=False): """A decorator which can be applied to any function which returns a :py:class:`.StructureSet` and which takes no other parameters 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) original = {s: n for n, s in enumerate(objects.structures)} if len(args) == 1: return {objects.get(args[0])} if args[0] in objects.ids else set() for k, v in kwargs.items(): objects = filter_objects(objects, k, v) if tuple_: return tuple(sorted( objects.structures, key=lambda s: original[s] )) else: return set(objects.structures) return structures
[docs]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. In atomium, various classes define methods like atoms, residues, etc. - this decorator can make a function like atom, residue which takes all the same params but just returns one object. :param func: the function to modify. :rtype: ``function``""" def structure(self, *args, **kwargs): for obj in func(self, *args, **kwargs): return obj return structure
[docs]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
[docs]class StructureSet: """A data structure for holding structures. It stores them internally as a dictionary where they keys are IDs (to allow rapid lookup by ID) and the values are all structures with that ID (to allow for duplicate IDs). Two structure sets can be added together, but they are immutable - the structures they have when they are made is the structures they will always have. They're basically sets optimised to lookup things by ID. :param \* args: the structures that will make up the StructureSet.""" def __init__(self, *args): self._d = {} for obj in args: if obj._id in self._d: self._d[obj._id].add(obj) else: self._d[obj._id] = {obj} def __add__(self, other): new = StructureSet() for s in (self, other): for key, value in s._d.items(): if key in new._d: new._d[key].update(value) else: new._d[key] = value return new def __len__(self): return len(self.structures) @property def ids(self): """Returns the IDs of the StructureSet. :rtype: ``set``""" return self._d.keys() @property def structures(self): """Returns the structures of the StructureSet. :rtype: ``list``""" structures = [] for s in self._d.values(): structures += s return structures
[docs] def get(self, id): """Gets a structure by ID. If an ID points to multiple structures, just one will be returned. :returns: some structure.""" matches = self._d.get(id, set()) for match in matches: return match