Source code for schrodinger.protein.seqres

"""
Functions for getting and setting SEQRES data on a structure.
"""

import decorator

from schrodinger.infra import mm
from schrodinger.infra.util import OneIndexedList

_SEQRES_BLOCK = "m_PDB_SEQRES"
_SEQRES_PROPS = ["s_pdb_chain_id", "s_pdb_SEQRES"]


@decorator.decorator
def _mminit(func, *args, **kwargs):
    """
    A decorator for initializing and terminating the mmlibs required for reading
    and writing SEQRES data.
    """
    try:
        mm.mmerr_initialize()
        mm.m2io_initialize(mm.error_handler)
        mm.mmct_initialize(mm.error_handler)
        return func(*args, **kwargs)
    finally:
        mm.mmct_terminate()
        mm.m2io_terminate()
        mm.mmerr_terminate()


@_mminit
def get_seqres(st):
    """
    Read in all SEQRES data from the structure.

    :param st: The structure to read SEQRES data from
    :type st: schrodinger.structure.Structure

    :return: A dictionary of {chain name: list of residue names} (ordered
        identically to the SEQRES data in the `Structure`) or None if there's no
        SEQRES information present in the structure.  Note that the lists of
        residue names uses a OneIndexedList since the i_pdb_seqres_index residue
        property values assume a starting index of one.
    :rtype: dict(str, OneIndexedList[str]) or NoneType
    """
    try:
        data_handle = mm.mmct_ct_m2io_get_unrequested_handle(st)
    except mm.MmException:
        data_handle = mm.mmct_ct_get_or_open_additional_data(st, True)
    num_seqres_blocks = mm.m2io_get_number_blocks(data_handle, _SEQRES_BLOCK)
    if not num_seqres_blocks:
        # There's no seqres information present
        return None
    mm.m2io_goto_block(data_handle, _SEQRES_BLOCK, 1)
    try:
        num_rows = mm.m2io_get_index_dimension(data_handle)
        seqres_by_chain = {}
        for row in range(1, num_rows + 1):
            chain, seqres = mm.m2io_get_string_indexed(data_handle, row,
                                                       _SEQRES_PROPS)
            seqres = OneIndexedList(seqres.split())
            seqres_by_chain[chain] = seqres
    finally:
        mm.m2io_leave_block(data_handle)
    return seqres_by_chain


@_mminit
def set_seqres(st, seqres):
    """
    Add SEQRES data to a structure.  Any SEQRES data that was previously present
    in the structure will be overwritten.

    :param st: The structure to add SEQRES data to.
    :type st: schrodinger.structure.Structure

    :param seqres: A dictionary of {chain name: list of residue names}.  May be
        None or an empty dictionary to clear SEQRES data.
    :type seqres: dict(str, list[str]) or None
    """
    data_handle = mm.mmct_ct_get_or_open_additional_data(st, True)
    # if there's any SEQRES data present, delete it (otherwise the new data will
    # be ignored)
    if mm.m2io_get_number_blocks(data_handle, _SEQRES_BLOCK):
        mm.m2io_delete_named_block(data_handle, _SEQRES_BLOCK)
    if not seqres:
        # There's no new data to write
        return
    mm.m2io_open_block(data_handle, _SEQRES_BLOCK)
    try:
        mm.m2io_set_index_dimension(data_handle, len(seqres))
        for idx, (chain_name, res_names) in enumerate(seqres.items(), start=1):
            res_names = " ".join(res_names)
            mm.m2io_put_string_indexed(data_handle, idx, _SEQRES_PROPS,
                                       [chain_name, res_names])
    finally:
        mm.m2io_close_block(data_handle)


@_mminit
def has_seqres(st):
    """
    Determine if a structure contains SEQRES data.

    :param st: The structure to read SEQRES data from.
    :type st: schrodinger.structure.Structure

    :return: Whether the structure contains SEQRES data.
    :rtype: bool
    """
    try:
        data_handle = mm.mmct_ct_m2io_get_unrequested_handle(st)
    except mm.MmException:
        data_handle = mm.mmct_ct_get_or_open_additional_data(st, True)
    num_seqres_blocks = mm.m2io_get_number_blocks(data_handle, _SEQRES_BLOCK)
    return num_seqres_blocks > 0