diff --git a/src/graph_scheduler/__init__.py b/src/graph_scheduler/__init__.py index 91d0b1b..0b73b54 100644 --- a/src/graph_scheduler/__init__.py +++ b/src/graph_scheduler/__init__.py @@ -11,6 +11,7 @@ Documentation: https://kmantel.github.io/graph-scheduler/ """ +import inspect import pint _unit_registry = pint.get_application_registry() @@ -22,10 +23,24 @@ from . import time # noqa: E402 from . import utilities # noqa: E402 -from .condition import * # noqa: E402 -from .scheduler import * # noqa: E402 -from .time import * # noqa: E402 -from .utilities import * # noqa: E402 + +condition.__all__ = [] +for k, v in condition.__dict__.items(): + if ( + ( + k[0] != '_' + and inspect.isclass(v) + and issubclass(v, condition.Condition) + ) + or k in condition._additional__all__ + ): + condition.__all__.append(k) + + +from .condition import * # noqa: E402, F401, F403 +from .scheduler import * # noqa: E402, F401, F403 +from .time import * # noqa: E402, F401, F403 +from .utilities import * # noqa: E402, F401, F403 __all__ = list(condition.__all__) __all__.extend(scheduler.__all__) @@ -38,4 +53,5 @@ from . import _version # noqa: E402 __version__ = _version.get_versions()['version'] +del inspect del pint diff --git a/src/graph_scheduler/condition.py b/src/graph_scheduler/condition.py index aa6f4ef..5247783 100644 --- a/src/graph_scheduler/condition.py +++ b/src/graph_scheduler/condition.py @@ -298,14 +298,8 @@ def converge(node, thresh): from graph_scheduler.time import TimeScale from graph_scheduler.utilities import call_with_pruned_args -__all__ = [ - 'AfterCall', 'AfterNCalls', 'AfterNCallsCombined', 'AfterNPasses', 'AfterNConsiderationSetExecutions', 'AfterNEnvironmentStateUpdates', 'AfterPass', - 'AtEnvironmentSequence', 'AfterEnvironmentSequence', 'AfterNEnvironmentSequences', 'AfterConsiderationSetExecution', 'AfterEnvironmentStateUpdate', 'All', 'AllHaveRun', 'Always', 'And', 'Any', - 'AtNCalls','AtPass', 'AtEnvironmentSequenceStart', 'AtEnvironmentSequenceNStart', 'AtConsiderationSetExecution', 'AtEnvironmentStateUpdate', - 'AtEnvironmentStateUpdateStart', 'AtEnvironmentStateUpdateNStart', 'BeforeNCalls', 'BeforePass', 'BeforeConsiderationSetExecution', 'BeforeEnvironmentStateUpdate', - 'Condition','ConditionError', 'ConditionSet', 'EveryNCalls', 'EveryNPasses', - 'JustRan', 'Never', 'Not', 'NWhen', 'Or', 'WhenFinished', 'WhenFinishedAll', 'WhenFinishedAny', 'While', 'WhileNot', 'TimeInterval', 'TimeTermination', 'Threshold' -] + +_additional__all__ = ['ConditionError', 'ConditionSet'] logger = logging.getLogger(__name__)