From 3cfc0db136b652d7b4515a7ecc78ca4b8db8dcd8 Mon Sep 17 00:00:00 2001
From: Alexander Hess <alexander@webartifex.biz>
Date: Wed, 18 Sep 2024 18:29:40 +0200
Subject: [PATCH] Organize top-level imports for `lalib.elements`

- make `GF2`, `one`, and `zero`, defined in the `lalib.elements.gf2`
  module, available as top-level imports in the `lalib.elements`
  sub-package via `from lalib.elements import *`
- provide some code snippets in the sub-package's docstring
- test the star import
---
 src/lalib/elements/__init__.py           | 35 +++++++++++++++++++++++-
 tests/elements/test_top_level_imports.py | 18 ++++++++++++
 tests/test_docstrings.py                 |  1 +
 3 files changed, 53 insertions(+), 1 deletion(-)
 create mode 100644 tests/elements/test_top_level_imports.py

diff --git a/src/lalib/elements/__init__.py b/src/lalib/elements/__init__.py
index ae48619..fe5dc37 100644
--- a/src/lalib/elements/__init__.py
+++ b/src/lalib/elements/__init__.py
@@ -1 +1,34 @@
-"""Various elements of various fields."""
+"""Various elements of various fields.
+
+Import the objects like so:
+
+>>> from lalib.elements import *
+
+Then, use them:
+
+>>> one + zero
+one
+
+>>> GF2(0)
+zero
+>>> GF2(42)
+Traceback (most recent call last):
+...
+ValueError: ...
+>>> GF2(42, strict=False)
+one
+"""
+
+from lalib.elements import gf2
+
+
+GF2, one, zero = gf2.GF2, gf2.one, gf2.zero
+
+del gf2
+
+
+__all__ = (
+    "GF2",
+    "one",
+    "zero",
+)
diff --git a/tests/elements/test_top_level_imports.py b/tests/elements/test_top_level_imports.py
new file mode 100644
index 0000000..dd19bcd
--- /dev/null
+++ b/tests/elements/test_top_level_imports.py
@@ -0,0 +1,18 @@
+"""Test top-level imports for `lalib.elements`."""
+
+from lalib import elements as top_level
+
+
+def test_top_level_imports():
+    """Verify `from lalib.elements import *` works."""
+    environment = {}
+
+    exec("...", environment, environment)  # noqa: S102
+    defined_vars_before = set(environment)
+
+    exec("from lalib.elements import *", environment, environment)  # noqa: S102
+    defined_vars_after = set(environment)
+
+    new_vars = defined_vars_after - defined_vars_before
+
+    assert new_vars == set(top_level.__all__)
diff --git a/tests/test_docstrings.py b/tests/test_docstrings.py
index 89ae74d..dd4ac9f 100644
--- a/tests/test_docstrings.py
+++ b/tests/test_docstrings.py
@@ -14,6 +14,7 @@ import xdoctest
     "module",
     [
         "lalib",
+        "lalib.elements",
         "lalib.elements.gf2",
     ],
 )