diff --git a/pyflow/attributes.py b/pyflow/attributes.py index 559ceb4..6ef5947 100644 --- a/pyflow/attributes.py +++ b/pyflow/attributes.py @@ -939,6 +939,10 @@ class InLimit(Attribute): Parameters: value(str,Limit_): The name of the limit or a limit object. + path(str): The optional path to the limit if the limit is not in the same node as the InLimit attribute. + tokens(int): The number of tokens to consume from the limit when a task is submitted. + limit_this_node_only(bool): Whether the limit should only apply to current node. + limit_submission(bool): Whether the limit should only apply to submissions Example:: @@ -946,8 +950,19 @@ class InLimit(Attribute): pyflow.InLimit(l) """ - def __init__(self, value): + def __init__( + self, + value: str | Limit, + path: str = "", + tokens: int = 1, + limit_this_node_only: bool = False, + limit_submission: bool = False, + ): super().__init__("_" + str(value), value) + self.path = path + self.tokens = tokens + self.limit_this_node_only = limit_this_node_only + self.limit_submission = limit_submission def _build(self, ecflow_parent): value = self.value @@ -955,9 +970,32 @@ def _build(self, ecflow_parent): return if isinstance(value, Limit): value = value.fullname.split(":") - ecflow_parent.add_inlimit(ecflow.InLimit(value[1], value[0])) + if self.path: + if self.path != value[0]: + raise ValueError( + "InLimit path {} does not match limit path {}".format( + self.path, value[0] + ) + ) + ecflow_parent.add_inlimit( + ecflow.InLimit( + value[1], + value[0], + self.tokens, + self.limit_this_node_only, + self.limit_submission, + ) + ) else: - ecflow_parent.add_inlimit(ecflow.InLimit(str(value))) + ecflow_parent.add_inlimit( + ecflow.InLimit( + str(value), + self.path, + self.tokens, + self.limit_this_node_only, + self.limit_submission, + ) + ) class Inlimit(InLimit): diff --git a/tests/test_inlimits.py b/tests/test_inlimits.py index a0e9f7e..a4866fc 100644 --- a/tests/test_inlimits.py +++ b/tests/test_inlimits.py @@ -1,4 +1,6 @@ -from pyflow import Limits, Suite, Tasks +import pytest + +from pyflow import InLimit, Limit, Limits, Suite, Tasks def test_inlimits(): @@ -14,6 +16,29 @@ def test_inlimits(): s.generate_node() +@pytest.mark.parametrize( + "options", + [ + {}, + {"path": "/s"}, + {"value": "tlimit"}, + {"value": "tlimit", "path": "/s"}, + {"value": "tlimit", "tokens": 1}, + {"value": "tlimit", "limit_this_node_only": True}, + {"value": "tlimit", "limit_submission": True}, + ], +) +def test_options(options): + with Suite("s") as s: + limit = Limit("tlimit", value=3) + if "value" not in options: + options["value"] = limit + Tasks("t", "t2", inlimits=InLimit(**options)) + + s.check_definition() + s.generate_node() + + if __name__ == "__main__": from os import path