diff --git a/snowfakery/template_funcs.py b/snowfakery/template_funcs.py index df751d97..47e985f6 100644 --- a/snowfakery/template_funcs.py +++ b/snowfakery/template_funcs.py @@ -144,8 +144,20 @@ def random_number(self, min: int, max: int, step: int = 1) -> int: """Pick a random number between min and max like Python's randint.""" return random.randrange(min, max + 1, step) - def reference(self, x: Any): + def reference( + self, x: Any = None, object: str = None, id: Union[str, int] = None + ): """YAML-embeddable function to Reference another object.""" + if x is not None: + return self._reference_from_scalar(x) + elif object: + try: + id = int(id) + except TypeError: + raise DataGenError("Cannot interpret id as integer") + return ObjectReference(object, id) + + def _reference_from_scalar(self, x: Any): if hasattr(x, "id"): # reference to an object with an id target = x elif isinstance(x, str): # name of an object diff --git a/tests/test_references.py b/tests/test_references.py index 24574eb1..d7c6af55 100644 --- a/tests/test_references.py +++ b/tests/test_references.py @@ -518,3 +518,23 @@ def test_random_reference_to_nickname_fails(self): with pytest.raises(DataGenError) as e: generate(StringIO(yaml)) assert "there is no table named parent" in str(e).lower() + + def test_reference_by_id(self, generated_rows): + yaml = """ + - object: Parent + nickname: ParentNickname + just_once: true + + - object: Child + fields: + parent1: + reference: ParentNickname + parent2: + reference: + object: Parent + id: 1 + """ + + generate(StringIO(yaml)) + child = generated_rows.table_values("Child", 0) + assert child["parent1"] == child["parent2"]