1#
2# This file is part of the GROMACS molecular simulation package.
3#
4# Copyright (c) 2019, by the GROMACS development team, led by
5# Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
6# and including many others, as listed in the AUTHORS file in the
7# top-level source directory and at http://www.gromacs.org.
8#
9# GROMACS is free software; you can redistribute it and/or
10# modify it under the terms of the GNU Lesser General Public License
11# as published by the Free Software Foundation; either version 2.1
12# of the License, or (at your option) any later version.
13#
14# GROMACS is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
17# Lesser General Public License for more details.
18#
19# You should have received a copy of the GNU Lesser General Public
20# License along with GROMACS; if not, see
21# http://www.gnu.org/licenses, or write to the Free Software Foundation,
22# Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
23#
24# If you want to redistribute modifications to GROMACS, please
25# consider that scientific software is very special. Version
26# control is crucial - bugs must be traceable. We will be happy to
27# consider code for inclusion in the official distribution, but
28# derived work must not be called official GROMACS. Details are found
29# in the README & COPYING files - if they are missing, get the
30# official version at http://www.gromacs.org.
31#
32# To help us fund GROMACS development, we humbly ask that you cite
33# the research papers on the package. Check out http://www.gromacs.org.
34
35"""Extra support for type hinting and generics.
36
37Provides additional support for annotation and static type checking.
38Extends the specifications of the abstract base classes in gmxapi.abc.
39"""
40
41from typing import TypeVar, Generic, NewType, Type
42
43import gmxapi.abc
44
45# Use SourceTypeVar and ResultTypeVar for static type hints, annotations, and as a parameter to generics.
46# Use valid_source_types and valid_result_types for run-time type checking.
47ResultTypeVar = TypeVar('ResultTypeVar', *(str, bool, int, float, dict, gmxapi.abc.NDArray))
48valid_result_types = ResultTypeVar.__constraints__
49
50SourceTypeVar = TypeVar('SourceTypeVar',
51                        *(str, bool, int, float, dict, gmxapi.abc.NDArray, gmxapi.abc.EnsembleDataSource))
52valid_source_types = SourceTypeVar.__constraints__
53
54# Place holder for type annotations of Context objects.
55# TODO: Expand to support some static type checking.
56_Context = TypeVar('_Context', bound=gmxapi.abc.Context)
57# Type variable that binds to subclasses of the forward referenced OperationImplentation ABC.
58_Op = TypeVar('_Op', bound=gmxapi.abc.OperationImplementation)
59
60# Note: We could make many types generic in SourceContext and TargetContext and
61# compose functionality into instances of a helpful base class with the option of
62# subclassing the generic base with concrete types. If we want to reserve the word
63# "Context" for an execution context, then we can say that a SourceResource or
64# TargetEnvironment can each be either a Context or an OperationExecutionDirector.
65#
66# We can't
67# generally add type-hinting details in terms of a specialized generic, but we
68# can require a particular version of a generic and we can use the inferred
69# type parameter value in limited cases, including to fully specify other generics.
70# For instance, we can't assert that a generic with a particular parameter has
71# additional attributes, but we can overload a function that produces an object
72# with those attributes. More generally, using some sort of identity function like
73# that may be the solution to the type erasure of a generic abstract base class.
74# For static type checking of dynamically defined subclasses, we can use a decorator
75# to reduce boiler plate if the type checker does a decent job of resolving closures.
76# Ironically, there may be cases where we want to *avoid* annotating the return
77# value to allow the static source analyzer to infer the dynamically typed return value.
78
79
80class Future(Generic[ResultTypeVar]):
81    """Generic result data."""
82    def dtype(self) -> Type[ResultTypeVar]:
83        ...
84
85    def result(self) -> ResultTypeVar:
86        ...
87
88
89class OperationReference(gmxapi.abc.OperationReference, Generic[_Op]):
90    """Object with an OperationReference interface.
91
92    Generic version of AbstractOperationReference, parameterized by operation
93    implementation.
94    """
95
96
97class OperationDirector(Generic[_Op, _Context]):
98    """Generic abstract operation director.
99
100    An operation director is instantiated for a specific operation and context
101    by a dispatching factory to add a computational element to the work managed
102    by the context.
103
104    Note:
105        This is a generic class, as defined with the Python `typing` module.
106        If used as a base class, re-expression of the TypeVar parameters in class
107        subscripts will cause the derived class to be generic unless regular
108        classes (non-TypeVar) are given. Omission of the subscripts causes `Any`
109        to be bound, which also results in a non-generic subclass.
110    """
111
112
113class OperationImplementation(Generic[_Op], gmxapi.abc.OperationImplementation):
114    """Essential interface of an Operation implementation.
115
116    Describe the essential features of an Operation that can be registered with
117    gmxapi to support building and executing work graphs in gmxapi compatible
118    execution contexts.
119    """
120    # The executable part of an operation consumes a distinct resource type.
121    # The resource type may be opaque, because it is created through a factory
122    # and only used in passing to a function object.
123    ResourceType = NewType('ResourceType', object)
124
125    def name(self) -> str:
126        """The name of the operation.
127
128        Generally, this corresponds to a callable attribute of a Python module
129        (named by namespace()) that acts as a factory for new operation instances.
130        It is also used by Context implementations to locate code supporting
131        the operation.
132        """
133        ...
134
135    def namespace(self) -> str:
136        """The namespace of the operation.
137
138        Generally, the namespace corresponds to a Python module importable in
139        the execution environment.
140        """
141
142    # Note that this indicates that the ABC requires subclasses to provide a generic function,
143    # which _is_ the factory behavior we are trying to specify.
144    # TODO: Confirm that this type-checks correctly.
145    # We may need to look more at how the type checking is implemented to see how
146    # to do this right.
147    def director(self, context: _Context) -> OperationDirector[_Op, _Context]:
148        """Factory to get an OperationDirector appropriate for the context."""
149        ...
150