From b71f7a2759bbdd3a61bd0701eb568ac12d27fcf3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Timoth=C3=A9e=20Mazzucotelli?= <pawamoy@pm.me>
Date: Wed, 12 Jan 2022 13:58:55 +0100
Subject: [PATCH] tests: Add helpers to create object trees

---
 tests/helpers.py | 70 +++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 69 insertions(+), 1 deletion(-)

diff --git a/tests/helpers.py b/tests/helpers.py
index 7818f29f..8cb5b2b5 100644
--- a/tests/helpers.py
+++ b/tests/helpers.py
@@ -8,7 +8,7 @@
 from typing import Iterator
 
 from griffe.agents.inspector import inspect
-from griffe.dataclasses import Module
+from griffe.dataclasses import Module, Object
 from tests import TESTS_DIR, TMP_DIR
 
 
@@ -42,3 +42,71 @@ def temporary_inspected_module(code: str) -> Iterator[Module]:
     """
     with temporary_pyfile(code) as (name, path):
         yield inspect(name, filepath=path)
+
+
+def vtree(*objects: Object, return_leaf: bool = False) -> Object:
+    """Link objects together, vertically.
+
+    Parameters:
+        *objects: A sequence of objects. The first one is at the top of the tree.
+        return_leaf: Whether to return the leaf instead of the root.
+
+    Raises:
+        ValueError: When no objects are provided.
+
+    Returns:
+        The top or leaf object.
+    """
+    if not objects:
+        raise ValueError("At least one object must be provided")
+    top = objects[0]
+    leaf = top
+    for obj in objects[1:]:
+        leaf[obj.name] = obj
+        leaf = obj
+    return leaf if return_leaf else top
+
+
+def htree(*objects: Object) -> Object:
+    """Link objects together, horizontally.
+
+    Parameters:
+        *objects: A sequence of objects. All objects starting at the second become members of the first.
+
+    Raises:
+        ValueError: When no objects are provided.
+
+    Returns:
+        The first given object, with all the other objects as members of it.
+    """
+    if not objects:
+        raise ValueError("At least one object must be provided")
+    top = objects[0]
+    for obj in objects[1:]:
+        top[obj.name] = obj
+    return top
+
+
+def module_vtree(path, leaf_package: bool = True, return_leaf: bool = False) -> Module:
+    """Link objects together, vertically.
+
+    Parameters:
+        path: The complete module path, like `"a.b.c.d"`.
+        leaf_package: Whether the deepest module should also be a package.
+        return_leaf: Whether to return the leaf instead of the root.
+
+    Raises:
+        ValueError: When no objects are provided.
+
+    Returns:
+        The top or leaf module.
+    """
+    parts = path.split(".")
+    modules = [Module(name, filepath=Path(*parts[:index], "__init__.py")) for index, name in enumerate(parts)]
+    if not leaf_package:
+        try:
+            filepath = modules[-1].filepath.with_stem(parts[-1])  # type: ignore[attr-defined]
+        except AttributeError:  # TODO: remove once Python 3.8 is dropped
+            filepath = modules[-1].filepath.with_name(f"{parts[-1]}.py")
+        modules[-1]._filepath = filepath  # noqa: WPS437
+    return vtree(*modules, return_leaf=return_leaf)  # type: ignore[return-value]