Source code for schrodinger.infra.mmim

"""
Wrapper classes and functions for dealing with the MMIM library, which is
used to read & write Impact and QSite input files.

See MMIMDict documentation for more details.

Copyright Schrodinger, LLC. All rights reserved.

"""

from collections.abc import MutableMapping

import schrodinger.infra.mm as mm
from schrodinger.infra import mmkv
from schrodinger.infra.mmkv import MMKVArgList

#
# MMIM stores a large enumerated list of keywords by data type, and these
# are surrounded by MMIM_ARG_*_BEGIN and MMIM_ARG_*_END values. (See the
# mmlibs/mmim/mmim_def.h header file if this doesn't make sense to you.)
# In order for this module to call the right functions in setting and
# getting keyword values, the following dict indexes the various functions
# by the MMIM_ARG_*_END value.
#
# This dict stores a tuple of
#  (
#   beginning MMIM index for the arg type,
#   string representation of the type,
#   converter function,
#   getter function,
#   setter function,
#   array size function,
#   array size set function,
#  )
#
# The array size function is None if the type isn't an array type.
#

_mmim_arg_handlers = {
    mm.MMIM_ARG_BOOL_END: (
        mm.MMIM_ARG_BOOL_BEGIN,
        "boolean",
        mmkv._convert_to_bool,
        mm.mmim_handle_get_arg_bool,
        mm.mmim_handle_set_arg_bool,
        None,
        None,
    ),
    mm.MMIM_ARG_BOOL_ARRAY_END: (
        mm.MMIM_ARG_BOOL_ARRAY_BEGIN,
        "boolean array",
        mmkv._convert_to_bool_array,
        mm.mmim_handle_get_indexed_arg_bool,
        mm.mmim_handle_set_indexed_arg_bool,
        mm.mmim_handle_get_indexed_arg_bool_count,
        mm.mmim_handle_set_indexed_arg_bool_count,
    ),
    mm.MMIM_ARG_FLOAT_END: (
        mm.MMIM_ARG_FLOAT_BEGIN,
        "float",
        float,
        mm.mmim_handle_get_arg_float,
        mm.mmim_handle_set_arg_float,
        None,
        None,
    ),
    mm.MMIM_ARG_FLOAT_ARRAY_END: (
        mm.MMIM_ARG_FLOAT_ARRAY_BEGIN,
        "float array",
        mmkv._convert_to_float_array,
        mm.mmim_handle_get_indexed_arg_float,
        mm.mmim_handle_set_indexed_arg_float,
        mm.mmim_handle_get_indexed_arg_float_count,
        mm.mmim_handle_set_indexed_arg_float_count,
    ),
    mm.MMIM_ARG_INT_END: (
        mm.MMIM_ARG_INT_BEGIN,
        "integer",
        int,
        mm.mmim_handle_get_arg_int,
        mm.mmim_handle_set_arg_int,
        None,
        None,
    ),
    mm.MMIM_ARG_INT_ARRAY_END: (
        mm.MMIM_ARG_INT_ARRAY_BEGIN,
        "integer array",
        mmkv._convert_to_int_array,
        mm.mmim_handle_get_indexed_arg_int,
        mm.mmim_handle_set_indexed_arg_int,
        mm.mmim_handle_get_indexed_arg_int_count,
        mm.mmim_handle_set_indexed_arg_int_count,
    ),
    mm.MMIM_ARG_STRING_END: (
        mm.MMIM_ARG_STRING_BEGIN,
        "string",
        lambda x: x,  # Return without modifying
        mm.mmim_handle_get_arg_string,
        mm.mmim_handle_set_arg_string,
        None,
        None,
    ),
    mm.MMIM_ARG_STRING_ARRAY_END: (
        mm.MMIM_ARG_STRING_ARRAY_BEGIN,
        "string array",
        mmkv._convert_to_string_array,
        mm.mmim_handle_get_indexed_arg_string,
        mm.mmim_handle_set_indexed_arg_string,
        mm.mmim_handle_get_indexed_arg_string_count,
        mm.mmim_handle_set_indexed_arg_string_count,
    ),
}


[docs]def convert_string(key, value): """ Given a key (as integer or string) and a value (string), try to convert the value to the proper type for this keyword. Return converted value. Will raise a KeyError if the wrong key value is specified. NOTE: Used by glide.py """ if not isinstance(key, int): key = mm.mmim_get_index_by_name(key.lower()) handlers = mmkv.get_handlers_for_key(key, _mmim_arg_handlers) converter = handlers[2] return converter(value)
[docs]class MMIMArgList(MMKVArgList):
[docs] def __init__(self, handle, key, converter=None, getter=None, setter=None, get_len=None, set_len=None): """ Create an instance from the MMKV/MMIM handle and key (int or string). Takes functions to get and set indexed values, and a function to determine the length. If these are not provided, they are looked up based on the key provided. get_len and set_len methods should be defined for array types only. """ # FIXME add ability to initialize from a Python list of strings, ints, # floats, or booleans. self.handle = handle self.key = key # Allow getter, setter, list_len to be provided at creation time, # but look them up automatically if they aren't. if getter is None: _, _, converter, getter, setter, get_len, set_len = mmkv.get_handlers_for_key( key, _mmim_arg_handlers) self._converter = converter self._getter = getter self._setter = setter self._get_len = get_len self._set_len = set_len
[docs]class MMIMDict(MutableMapping): """ A class for setting and getting MMIM keyword values. Given an MMIM handle, create an instance as follows: mydict = MMIMDict(mmim_handle) Then set any keywords: mydict[<keywordname>] = <value> """ _NOT_IMPLEMENTED_MSG = "mmim cannot report known keys."
[docs] def __init__(self, handle): """ Create an instance from the given MMIM handle. """ self.handle = handle # Cache the list objects so they don't have to be created on every # access and also so they maintain a constant id. self._list_objs = {}
def __getitem__(self, key): """ Return the value stored in the MMIM handle for the given key (integer or string). Note that if no explicit setting has been made for a keyword that the default value will be returned. """ if key in self._list_objs: return self._list_objs[key] key, converter, getter, setter, get_len, set_len = self._get_handlers_for_key( key) try: # If there's no get_len method, then the key data type is a # scalar and the simple getter can be called. if get_len is None: return getter(self.handle, key) else: # If the data type is a list type, return an MMIMArgList # instance. list_obj = MMIMArgList(self.handle, key, converter, getter, setter, get_len, set_len) self._list_objs[key] = list_obj return list_obj except mm.MmException as e: if e.rc == mm.MMIM_ERR: raise KeyError("MMIM key %d undefined for handle %d" % (key, self.handle)) def __setitem__(self, key, value): """ Set the value for the provided key (integer or string) to <value>. If the key type is an MMIM arglist (array), then the value provided must be an list with elements of the proper type. """ key, converter, getter, setter, get_len, set_len = self._get_handlers_for_key( key) # If there's no get_len method, then the key data type is a scalar # and the simple setter can be called. if get_len is None: setter(self.handle, key, value) else: # If the data type is a list type # 1) check for an MMIMArgList in the cache, # 2) make sure it's the right size, # 3) copy the list values over into the MMIMArgList list_obj = self._list_objs.get(key) if list_obj is None: list_obj = MMIMArgList(self.handle, key, converter, getter, setter, get_len, set_len) self._list_objs[key] = list_obj # Initialize by setting the size: list_obj.setSize(len(value)) elif len(list_obj) != len(value): # Size has changed; adjust: list_obj.setSize(len(value)) for ix, new_value in enumerate(value): list_obj[ix] = new_value return def __delitem__(self, key): return mm.mmim_handle_unset_arg(self.handle, key)
[docs] def keys(self): # Without a way of implementing keys(), some dictionary functionality # will remain unavailable. # FIXME we may want to fix this raise NotImplementedError(self._NOT_IMPLEMENTED_MSG)
def _get_handlers_for_key(self, key): """ Retrieve the mmim functions to use in getting and setting the provided key (integer or string). Returns a tuple of key ID and converter/getter/setter functions. """ if not isinstance(key, int): key = mm.mmim_get_index_by_name(key.lower()) handlers = mmkv.get_handlers_for_key(key, _mmim_arg_handlers) return (key,) + handlers[2:]
[docs] def __len__(self): raise NotImplementedError(self._NOT_IMPLEMENTED_MSG)
def __iter__(self): raise NotImplementedError(self._NOT_IMPLEMENTED_MSG)