From 0393e969dd931ab91cd26c1bca552729403b5f5a Mon Sep 17 00:00:00 2001 From: Bas van Beek <43369155+BvB93@users.noreply.github.com> Date: Wed, 1 Dec 2021 12:45:58 +0100 Subject: [PATCH] Change `init` from a function into a context manager --- core/functions.py | 100 ++++++++++++++++++++-------- doc/source/components/functions.rst | 2 +- doc/source/conf.py | 2 +- 3 files changed, 73 insertions(+), 31 deletions(-) diff --git a/core/functions.py b/core/functions.py index 0478ec92..a7ddbb46 100644 --- a/core/functions.py +++ b/core/functions.py @@ -7,6 +7,7 @@ import time import types import warnings +import contextlib from typing import Callable, Dict, NoReturn from os.path import join as opj @@ -23,8 +24,7 @@ #=========================================================================== - -def init(path=None, folder=None, config_settings:Dict=None, quiet=False): +class init(contextlib.AbstractContextManager): """Initialize PLAMS environment. Create global ``config`` and the default |JobManager|. An empty |Settings| instance is created and populated with default settings by executing ``plams_defaults``. The following locations are used to search for the defaults file, in order of precedence: @@ -37,45 +37,84 @@ def init(path=None, folder=None, config_settings:Dict=None, quiet=False): Optionally, an additional `dict` (or |Settings| instance) can be provided to the `config_settings` argument which will be used to update the values from the ``plams_defaults``. + |init| can be either used as a standalone function or in conjunction with the ``with`` statement, the latter option automatically calling |finish| upon exiting the context manager: + + .. code-block:: python + + >>> from scm.plams import Molecule, Settings, AMSJob, init, finish + + >>> mol: Molecule = ... + >>> settings: Settings = ... + + >>> with init(): + ... job1 = AMSJob(molecule=mol, settings=settings) + ... result1 = job1.run() + + # Equivalently: + >>> init() + >>> job2 = AMSJob(molecule=mol, settings=settings) + >>> result2 = job2.run() + >>> finish() + + .. warning:: This function **must** be called before any other PLAMS command can be executed. Trying to do anything without it results in a crash. See also |master-script|. """ - if config.init: - return + def __init__(self, path=None, folder=None, config_settings:Dict=None, quiet=False, otherJM=None): + self.otherJM = otherJM - if 'PLAMSDEFAULTS' in os.environ and isfile(expandvars('$PLAMSDEFAULTS')): - defaults = expandvars('$PLAMSDEFAULTS') - elif 'AMSHOME' in os.environ and isfile(opj(expandvars('$AMSHOME'), 'scripting', 'scm', 'plams', 'plams_defaults')): - defaults = opj(expandvars('$AMSHOME'), 'scripting', 'scm', 'plams', 'plams_defaults') - else: - defaults = opj(dirname(dirname(__file__)), 'plams_defaults') - if not isfile(defaults): - raise PlamsError('plams_defaults not found, please set PLAMSDEFAULTS or AMSHOME in your environment') - with open(defaults, 'r') as f: - exec(compile(f.read(), defaults, 'exec')) + if config.init: + return - config.update(config_settings or {}) + if 'PLAMSDEFAULTS' in os.environ and isfile(expandvars('$PLAMSDEFAULTS')): + defaults = expandvars('$PLAMSDEFAULTS') + elif 'AMSHOME' in os.environ and isfile(opj(expandvars('$AMSHOME'), 'scripting', 'scm', 'plams', 'plams_defaults')): + defaults = opj(expandvars('$AMSHOME'), 'scripting', 'scm', 'plams', 'plams_defaults') + else: + defaults = opj(dirname(dirname(__file__)), 'plams_defaults') + if not isfile(defaults): + raise PlamsError('plams_defaults not found, please set PLAMSDEFAULTS or AMSHOME in your environment') + with open(defaults, 'r') as f: + exec(compile(f.read(), defaults, 'exec')) - from .jobmanager import JobManager - config.default_jobmanager = JobManager(config.jobmanager, path, folder) + config.update(config_settings or {}) - if not quiet: - log('Running PLAMS located in {}'.format(dirname(dirname(__file__))), 5) - log('Using Python {}.{}.{} located in {}'.format(*sys.version_info[:3], sys.executable), 5) - log('PLAMS defaults were loaded from {}'.format(defaults), 5) + from .jobmanager import JobManager + config.default_jobmanager = JobManager(config.jobmanager, path, folder) - log('PLAMS environment initialized', 5) - log('PLAMS working folder: {}'.format(config.default_jobmanager.workdir), 1) + if not quiet: + log('Running PLAMS located in {}'.format(dirname(dirname(__file__))), 5) + log('Using Python {}.{}.{} located in {}'.format(*sys.version_info[:3], sys.executable), 5) + log('PLAMS defaults were loaded from {}'.format(defaults), 5) - config.slurm = _init_slurm() if "SLURM_JOB_ID" in os.environ else None + log('PLAMS environment initialized', 5) + log('PLAMS working folder: {}'.format(config.default_jobmanager.workdir), 1) - try: - import dill - except ImportError: - log('WARNING: importing dill package failed. Falling back to the default pickle module. Expect problems with pickling', 1) + config.slurm = _init_slurm() if "SLURM_JOB_ID" in os.environ else None + + try: + import dill + except ImportError: + log('WARNING: importing dill package failed. Falling back to the default pickle module. Expect problems with pickling', 1) + + config.init = True - config.init = True + def __enter__(self): + """Enter the context manager; return |init|.""" + return self + + def __exit__(self, exc_type, exc_value, traceback): + """Exit the context manager; call |finish|.""" + finish(self.otherJM) + + def __repr__(self): + """Return :func:`repr(self) `.""" + return f"" + + def __bool__(self): + """Return :class:`bool(self) `.""" + return config.init def _init_slurm(): @@ -140,6 +179,9 @@ def finish(otherJM=None): This function must be called at the end of your script for |cleaning| to take place. See |master-script| for details. If you used some other job managers than just the default one, they need to be passed as *otherJM* list. + + This function is automatically called when using |init| with the ``with`` statement. + """ if not config.init: return diff --git a/doc/source/components/functions.rst b/doc/source/components/functions.rst index 5f23b7ed..c2647dc8 100644 --- a/doc/source/components/functions.rst +++ b/doc/source/components/functions.rst @@ -7,7 +7,7 @@ Public functions This chapter gathers information about public functions that can be used in PLAMS scripts. -.. autofunction:: init +.. autoclass:: init .. autofunction:: finish .. autofunction:: load .. autofunction:: load_all diff --git a/doc/source/conf.py b/doc/source/conf.py index 98bc636d..e76a803c 100644 --- a/doc/source/conf.py +++ b/doc/source/conf.py @@ -204,7 +204,7 @@ def setup(app): autodoc_typehints = 'none' rst_epilog = """ -.. |init| replace:: :func:`~scm.plams.core.functions.init` +.. |init| replace:: :class:`~scm.plams.core.functions.init` .. |log| replace:: :func:`~scm.plams.core.functions.log` .. |load| replace:: :func:`~scm.plams.core.functions.load` .. |load_all| replace:: :func:`~scm.plams.core.functions.load_all`