Source code for ase2sprkkr.asr.core.types

"""Implement custom types for ASR."""
import click
from ase.io import read
from ase.io.formats import UnknownFileTypeError
from asr.core import parse_dict_string


[docs] class AtomsFile(click.ParamType): """Read atoms object from filename and return Atoms object.""" name = "atomsfile"
[docs] def __init__(self, must_exist=True, *args, **kwargs): """Initialize AtomsFile object. Parameters ---------- must_exist : bool If False, errors relating to empty or missing files will be ignored and the returned atoms will be None in that case. If True, all errors will be raised if encountered. Returns ------- atoms : `ase.Atoms` """ self.must_exist = must_exist super().__init__(*args, **kwargs)
[docs] def convert(self, value, param, ctx): """Convert string to atoms object.""" try: return read(value, parallel=False) except (IOError, UnknownFileTypeError, StopIteration): if self.must_exist: raise return None
[docs] class DictStr(click.ParamType): """Read atoms object from filename and return Atoms object.""" name = "dictionary_string"
[docs] def convert(self, value, param, ctx): """Convert string to a dictionary.""" if isinstance(value, dict): return value return parse_dict_string(value)
[docs] def clickify_docstring(doc): """Take a standard docstring a make it Click compatible.""" if doc is None: return doc_n = doc.split('\n') clickdoc = [] skip = False for i, line in enumerate(doc_n): if skip: skip = False continue lspaces = len(line) - len(line.lstrip(' ')) spaces = ' ' * lspaces bb = spaces + '\b' if line.endswith('::'): skip = True if not doc_n[i - 1].strip(' '): clickdoc.pop(-1) clickdoc.extend([bb, line, bb]) else: clickdoc.extend([line, bb]) elif ('-' in line and (spaces + '-' * (len(line) - lspaces)) == line): clickdoc.insert(-1, bb) clickdoc.append(line) else: clickdoc.append(line) doc = '\n'.join(clickdoc) return doc