diff --git a/.github/workflows/pr-coverage.yml b/.github/workflows/pr-coverage.yml new file mode 100644 index 0000000..5617147 --- /dev/null +++ b/.github/workflows/pr-coverage.yml @@ -0,0 +1,103 @@ +name: PR Coverage + +on: + pull_request: + branches: + - main + +permissions: + contents: read + pull-requests: write + +env: + python-version: "3.11" + ALLOWED_COVERAGE_DECREASE: 0 + MIN_COVERAGE: 0.80 + MIN_DIFF_COVERAGE: 0.85 + +defaults: + run: + shell: bash + +jobs: + coverage: + runs-on: ubuntu-latest + outputs: + reference_percent_covered: ${{ steps.coverage_comment.outputs.reference_percent_covered }} + new_percent_covered: ${{ steps.coverage_comment.outputs.new_percent_covered }} + diff_total_percent_covered: ${{ steps.coverage_comment.outputs.diff_total_percent_covered }} + + steps: + - uses: actions/checkout@v5 + with: + fetch-depth: 0 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: ${{ env.python-version }} + + - name: Set up uv + uses: astral-sh/setup-uv@v6 + + - name: Get MAD binaries + run: | + mkdir -p ./src/pymadng/bin + curl https://madx.web.cern.ch/releases/madng/1.1/mad-linux-1.1.11 -o ./src/pymadng/bin/mad_Linux + curl https://madx.web.cern.ch/releases/madng/1.1/mad-macos-1.1.11 -o ./src/pymadng/bin/mad_Darwin + chmod +x ./src/pymadng/bin/mad_Linux ./src/pymadng/bin/mad_Darwin + + - name: Install package + run: uv sync --extra tfs --extra test + + - name: Run coverage + run: uv run pytest tests --cov=src/pymadng --cov-report=xml --cov-report=term + + - name: Coverage comment + id: coverage_comment + uses: py-cov-action/python-coverage-comment-action@v3 + with: + GITHUB_TOKEN: ${{ github.token }} + + check-coverage: + needs: coverage + runs-on: ubuntu-latest + steps: + - name: Check coverage decrease + env: + CURRENT_COVERAGE: ${{ needs.coverage.outputs.new_percent_covered }} + REF_COVERAGE: ${{ needs.coverage.outputs.reference_percent_covered }} + run: | + MIN_COVERAGE_DECR=$(echo "${REF_COVERAGE} - ${ALLOWED_COVERAGE_DECREASE}" | bc) + echo "Reference Coverage: ${REF_COVERAGE}" + echo "Current Coverage: ${CURRENT_COVERAGE}" + echo "Minimum Allowed Coverage: ${MIN_COVERAGE_DECR}" + if (( $(echo "${CURRENT_COVERAGE} < ${MIN_COVERAGE_DECR}" | bc -l) )); then + echo "Failed: coverage decreased (${CURRENT_COVERAGE} < ${MIN_COVERAGE_DECR})." + exit 1 + fi + echo "Success: coverage did not decrease." + + - name: Check total coverage + env: + CURRENT_COVERAGE: ${{ needs.coverage.outputs.new_percent_covered }} + run: | + echo "Current Coverage: ${CURRENT_COVERAGE}" + echo "Minimum Allowed Coverage: ${MIN_COVERAGE}" + if (( $(echo "${CURRENT_COVERAGE} < ${MIN_COVERAGE}" | bc -l) )); then + echo "Failed: coverage is below minimum (${CURRENT_COVERAGE} < ${MIN_COVERAGE})." + exit 1 + fi + echo "Success: coverage meets minimum requirement." + + - name: Check diff coverage + env: + CURRENT_DIFF_COVERAGE: ${{ needs.coverage.outputs.diff_total_percent_covered }} + run: | + echo "Current Diff Coverage: ${CURRENT_DIFF_COVERAGE}" + echo "Minimum Allowed Diff Coverage: ${MIN_DIFF_COVERAGE}" + if (( $(echo "${CURRENT_DIFF_COVERAGE} < ${MIN_DIFF_COVERAGE}" | bc -l) )); then + echo "Failed: diff coverage is below minimum (${CURRENT_DIFF_COVERAGE} < ${MIN_DIFF_COVERAGE})." + exit 1 + fi + echo "Success: diff coverage meets minimum requirement." diff --git a/.github/workflows/test-publish.yaml b/.github/workflows/test-publish.yaml index 339f797..e2f8150 100644 --- a/.github/workflows/test-publish.yaml +++ b/.github/workflows/test-publish.yaml @@ -12,7 +12,7 @@ jobs: fail-fast: false matrix: os: [ubuntu-latest, macos-latest] - python-version: ["3.11", "3.12", "3.13"] + python-version: ["3.11", "3.12", "3.13", "3.14"] steps: - uses: actions/checkout@v5 @@ -23,8 +23,8 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - python -m pip install pymadng[tfs] + python -m pip install pymadng[tfs] pytest - name: Test with python run: | - python -m unittest tests/*.py + python -m pytest tests diff --git a/.github/workflows/test-pymadng.yml b/.github/workflows/test-pymadng.yml index 076b5d3..6c78cde 100644 --- a/.github/workflows/test-pymadng.yml +++ b/.github/workflows/test-pymadng.yml @@ -22,7 +22,7 @@ jobs: fail-fast: false matrix: os: [ubuntu-latest, macos-latest] - python-version: ["3.11", "3.12", "3.13"] + python-version: ["3.11", "3.12", "3.13", "3.14"] steps: - uses: actions/checkout@v5 @@ -30,6 +30,8 @@ jobs: uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} + - name: Set up uv + uses: astral-sh/setup-uv@v6 - name: Get MAD Binaries run: | mkdir ./src/pymadng/bin @@ -38,8 +40,7 @@ jobs: chmod +x ./src/pymadng/bin/mad_Linux ./src/pymadng/bin/mad_Darwin - name: Install dependencies run: | - python -m pip install --upgrade pip - python -m pip install -e .[tfs] + uv sync --extra tfs --extra test - name: Test with python run: | - python -m unittest tests/*.py + uv run pytest tests diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..6ac2cf4 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,14 @@ +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.13.0 + hooks: + - id: ruff-check + args: [--fix] + - id: ruff-format + + - repo: https://github.com/streetsidesoftware/cspell-cli + rev: v9.3.3 + hooks: + - id: cspell + args: [--no-progress, --no-summary] + files: ^(README\.md|CHANGELOG\.md|docs/source/.*\.(md|rst))$ diff --git a/CHANGELOG.md b/CHANGELOG.md index 2f1f08e..39f8e50 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,9 @@ +0.10.0 (2026/03/11) \ +Major debug update: PyMAD-NG now exposes the MAD-NG debugger directly through `MAD.breakpoint()` and `MAD.pydbg()`, with MAD-side aliases (`breakpoint`, `pydbg`, and `python_breakpoint`) available inside executed code. \ +Debugger sessions now support scripted commands for tests and automation, plus improved interactive terminal handling with a stable Python-rendered prompt. \ +Quitting from the debugger now shuts down the current MAD session cleanly without leaving pipe state inconsistent. \ +Expanded debugger test coverage and documentation for Python-driven MAD debugging workflows. + 0.9.0 (2026/02/25) \ Remove support for Python 3.10, now only supporting Python 3.11 and above. \ Replaced incorrect Warning call with logging.warning for proper warning handling. \ diff --git a/README.md b/README.md index 5bc0ea0..6626c60 100644 --- a/README.md +++ b/README.md @@ -28,9 +28,16 @@ Before diving into PyMAD-NG, we recommend you: ### Explore Key Examples - **[LHC Matching Example](https://pymadng.readthedocs.io/en/latest/ex-lhc-couplingLocal.html)** – Real-world optics matching with intermediate feedback. -- **[Examples Page](https://pymadng.readthedocs.io/en/latest/examples.html)** - List of examples in an easy to read format. +- **[Examples Page](https://pymadng.readthedocs.io/en/latest/examples.html)** - List of examples in an easy to read format. - **[GitHub Examples Directory](https://github.com/MethodicalAcceleratorDesign/MAD-NG.py/blob/main/examples/)** – List of available examples on the repository +### A Few Things To Know Early + +- `MAD()` launches a real MAD-NG subprocess immediately. Prefer `with MAD() as mad:` in normal use. +- `mad["x"] = value` writes a variable inside MAD-NG. `mad.x = value` only changes the Python wrapper object. +- Communication is explicit: tell MAD-NG to `py:send(...)` before calling `mad.recv()`, and tell MAD-NG to `py:recv()` before sending Python data. +- Many high-level results are references to MAD-NG objects. Use `.eval()` or conversion helpers such as `.to_df()` when you want Python-side values. + If anything seems unclear: - Refer to the [API Reference](https://pymadng.readthedocs.io/en/latest/pymadng.html#module-pymadng) - Check the [MAD-NG Docs](https://madx.web.cern.ch/releases/madng/html/) @@ -59,7 +66,7 @@ Examples are stored in the `examples/` folder. Run any script with: ```bash -python3 examples/ex-fodos.py +python3 examples/ex-fodo/ex-fodos.py ``` You can also batch-run everything using: @@ -82,7 +89,7 @@ python3 runall.py ## 🤝 Contributing -We welcome contributions! See [`CONTRIBUTING.md`](docs/source/contributing.md) or the [Contributing Guide](https://pymadng.readthedocs.io/en/latest/contributing.html) in the docs. +We welcome contributions! See [the contributing guide](docs/source/contributing.md) or the [Contributing Guide](https://pymadng.readthedocs.io/en/latest/contributing.html) in the docs. Bug reports, feature requests, and pull requests are encouraged. diff --git a/cspell.config.yaml b/cspell.config.yaml new file mode 100644 index 0000000..d05dd3d --- /dev/null +++ b/cspell.config.yaml @@ -0,0 +1,68 @@ +version: "0.2" +language: en-GB + +ignorePaths: + - .git + - .venv + - build + - dist + +words: + - addopts + - assertf + - automodule + - beamline + - cmatrix + - CTPSA + - currentmodule + - dataframe + - dataframes + - defexpr + - eval-rst + - fodo + - fodos + - gmath + - gphys + - imatrix + - initialise + - kwargs + - ipairs + - irange + - lhcb + - linenos + - linters + - literalinclude + - logrange + - MADX + - MADNG + - madp + - matplotlib + - mtable + - mtbl + - ndarray + - numpy + - knl + - pathlib + - pyplot + - pydbg + - pymad + - pymadng + - PyPI + - pyproject + - pytest + - runall + - rtrn + - seqfile + - stdout + - stderr + - sterr + - testpaths + - tonumber + - tostring + - TPSA + - tpsa + - twiss + - venv + - vars + - xlabel + - ylabel diff --git a/docs/source/communication.md b/docs/source/communication.md index c5eb102..49fb8a9 100644 --- a/docs/source/communication.md +++ b/docs/source/communication.md @@ -16,6 +16,10 @@ Key points: - Data is retrieved via `{func}`MAD.recv`()` after explicit instruction to send it. - MAD-NG stdout is redirected to Python, but not intercepted. +```{tip} +Think of PyMAD-NG as controlling a persistent remote interpreter. Python and MAD-NG do not share memory; they exchange commands, references, and serialised data through pipes. +``` + ```{important} You must always **send instructions before sending data**, and **send a request before receiving data**. ``` @@ -34,6 +38,25 @@ mad.recv() # Receive the value → 42 Both {func}`MAD.send` and {func}`MAD.recv` are the core communication methods. See the {class}`pymadng.MAD` reference for more details. +### Two Common Patterns + +#### Pattern A: MAD-NG asks Python for data + +```python +mad.send("arr = py:recv()") +mad.send(my_array) +``` + +#### Pattern B: Python asks MAD-NG for data + +```python +mad.send("tbl = twiss {sequence=seq}") +mad.send("py:send(tbl)") +tbl = mad.recv("tbl") # Receives the table as a Python object that is interactive and can be converted to a DataFrame +``` + +If you keep these two patterns distinct, most send/receive bugs become much easier to reason about. + --- ## Supported Data Types @@ -63,14 +86,18 @@ The following types can be sent from Python to MAD-NG: - {func}`MAD.send` * - `start, stop, size` as float, int - `range`, `logrange` - - `mad.send_rng()`, `mad.send_lrng()` + - `mad.send_range()`, `mad.send_logrange()` * - Complex structures (e.g., TPSA, CTPSA) - `TPSA`, `CTPSA` - - `mad.send_tpsa()`, `mad.send_ctpsa()` + - `mad.send_tpsa()`, `mad.send_cpx_tpsa()` ``` For full compatibility, see the {mod}`pymadng.MAD` documentation. +```{note} +High-level MAD objects are often returned as references rather than copied Python objects. This is expected behaviour and is part of how PyMAD-NG avoids unnecessary data transfer. +``` + --- ## Converting TFS Tables to DataFrames @@ -113,6 +140,15 @@ mad.send(arr2) # DEADLOCK if previous data not yet received Always ensure each {func}`MAD.send` has a matching {func}`MAD.recv` if data is expected back. ``` +### A Safer Way to Think About It + +Before each operation, ask one question: + +- "Is MAD-NG waiting for Python?" +- or "Is Python waiting for MAD-NG?" + +If both sides are waiting to receive, the session will deadlock. + --- ## Scope: Local vs Global @@ -137,6 +173,65 @@ mad.send("print(a + (b or 5))") # b is nil → 10 + 5 = 15 Use `local` to avoid polluting the global MAD-NG namespace. ``` +### Python Assignment vs MAD Assignment + +This is worth stating explicitly because it is a common first-use mistake: + +```python +mad["x"] = 5 # writes x inside MAD-NG +mad.x = 5 # writes an attribute on the Python wrapper only +``` + +If you intend to create a MAD-NG variable, always use square-bracket assignment. + +--- + +## References and Materialising Values + +Many objects returned by the high-level interface are references to values living inside MAD-NG. + +```python +ref = mad.math.sin(1) +``` + +To materialise a Python value, if it exists (rather than a reference), use the `eval()` method: + +```python +value = ref.eval() +``` + +To materialise a table: + +```python +df = mad.tbl.to_df() +``` + +This distinction is especially important when inspecting objects interactively, writing assertions in tests, or passing results into ordinary Python libraries. + +--- + +## Executing Python Returned by MAD-NG + +{func}`MAD.recv_and_exec` executes Python code sent back from MAD-NG: + +```python +mad.send("py:send([[print('hello from MAD')]])") +mad.recv_and_exec() +``` + +The execution context automatically includes: + +- `mad`: the current PyMAD-NG session +- `breakpoint`: the debugger bridge +- `pydbg`: alias for the debugger bridge + +That means MAD-NG can request an interactive debugger stop like this: + +```python +mad.send("py:send([[breakpoint()]])") +mad.recv_and_exec() +``` + --- ## Customising the Environment @@ -174,10 +269,9 @@ See {meth}`pymadng.MAD.__init__` for all configuration options. ## Summary - Always match {func}`MAD.send` with {func}`MAD.recv` when data is expected. -- Use `mad.to_df()` for table conversion. +- Use `.to_df()` on MAD tables when you want a DataFrame. - Avoid deadlocks by receiving before sending again. - Manage scope using `local` wisely. - Use configuration flags to tailor behaviour. For more, see the {doc}`advanced_features`, {doc}`debugging`, and {doc}`function_reference` sections. - diff --git a/docs/source/contributing.md b/docs/source/contributing.md index c10748a..d94fe4f 100644 --- a/docs/source/contributing.md +++ b/docs/source/contributing.md @@ -22,18 +22,29 @@ This guide is for developers who want to contribute to PyMAD-NG or extend its ca ## Getting Started ### Prerequisites -- Python 3.7+ -- `numpy`, `pandas`, and optionally `tfs-pandas` +- Python 3.11+ +- `uv` for dependency and environment management - A valid MAD-NG executable (either system or bundled) -### Install PyMAD-NG in Editable Mode -To contribute to PyMAD-NG, clone the repository and install it in editable mode. This allows you to make changes and test them without reinstalling. -In specific cases, you may be allowed write access to the MAD-NG.py repository. If so, you only need to clone the repository and install it in editable mode. +### Set Up a Development Environment +Clone the repository, create the project environment with `uv`, and install the development tooling defined by the project extras. ```bash git clone https://github.com//MAD-NG.py.git -cd pymadng -pip install -e . +cd MAD-NG.py +uv sync --extra dev --extra test --extra tfs +``` + +That command creates the local `.venv` and installs: +- the package itself in editable mode +- test dependencies such as `pytest` and `pytest-cov` +- development tools such as `pre-commit` and `ruff` +- optional `tfs-pandas` support used by some tests and examples + +If you only need the core development tools, you can omit extras you do not need: + +```bash +uv sync --extra dev --extra test ``` --- @@ -47,15 +58,45 @@ git checkout -b your-feature-name 2. **Write Code** -3. **Test Code** - - All tests are located in the `tests/` directory, currently `unittests` are used. +3. **Run the Test Suite** +```bash +uv run pytest tests +``` + +4. **Run Coverage Locally** +```bash +uv run pytest tests --cov=src/pymadng --cov-report=term-missing +``` + +Coverage is configured in `pyproject.toml` to report on `src/pymadng` rather than the test files themselves. -4. **Run Tests** +5. **Run the linters** ```bash -python -m unittest tests/*.py +uv run ruff check . +uv run ruff format . ``` -5. **Submit a Pull Request** +6. **Enable and Run Pre-commit Hooks** +Install the hooks once per clone: + +```bash +uv run pre-commit install +``` + +Run all configured hooks manually: + +```bash +uv run pre-commit run --all-files +``` + +The repository currently uses: +- `ruff-check` for linting +- `ruff-format` for formatting +- `cspell` for spell-checking with British English (`en-GB`) configuration + +Spell-check configuration lives in `cspell.config.yaml`. + +7. **Submit a Pull Request** - Include a clear description of the change - Reference any related issues @@ -66,7 +107,14 @@ python -m unittest tests/*.py ### Code Style - Use descriptive names for everything - Keep high-level user APIs separate from internal helpers -- Use Ruff for code and import formatting. +- Use Ruff for linting and formatting. +- Prefer public API tests over mock-heavy tests when verifying subprocess behaviour. +- Update documentation whenever user-visible behaviour changes. + +### CI Expectations +- GitHub Actions uses `uv` to create the environment and run the test suite. +- Coverage is uploaded from the main Linux / Python 3.11 test job. +- Keep local commands aligned with CI when debugging failures. --- @@ -82,6 +130,6 @@ python -m unittest tests/*.py - Open a GitHub Issue - Tag maintainers in your Pull Request -- See the [Debugging Guide] or [Architecture Overview] for internals +- See the [Debugging Guide](debugging.md) or [Architecture Overview](architecture.md) for internals Thanks for contributing to PyMAD-NG! diff --git a/docs/source/debugging.md b/docs/source/debugging.md index 63d091d..9de0c05 100644 --- a/docs/source/debugging.md +++ b/docs/source/debugging.md @@ -41,7 +41,141 @@ This helps keep logs organised, especially when running long scripts. --- -## 2. Inspecting Command History +## 2. Interactive MAD Debugger + +PyMAD-NG now exposes MAD-NG's `MAD.dbg()` debugger directly through the Python wrapper. + +### 2.1 Entering the Debugger from Python + +Use {meth}`MAD.breakpoint` to stop inside MAD-NG and interact with the debugger from the Python terminal: + +```python +from pymadng import MAD + +with MAD() as mad: + mad.breakpoint() +``` + +There is also a shorter alias: + +```python +mad.pydbg() +``` + +When the Python process is attached to a real terminal, PyMAD-NG uses Python's line editor for the debugger prompt. That gives you normal command editing while still sending completed commands to `MAD.dbg()`. + +```{note} +The visible debugger prompt is re-drawn by Python, not by MAD-NG directly. This keeps cursor positioning, history navigation, and prompt colouring stable even though MAD-NG itself is running behind a pipe. +``` + +### 2.2 Entering the Debugger from MAD-NG Code + +During process startup, PyMAD-NG defines these MAD-side helper functions: + +- `python_breakpoint()` +- `pydbg()` +- `breakpoint()` + +That means ordinary MAD strings can drop into the debugger directly: + +```python +mad.send("breakpoint()") +``` + +or from within larger MAD programs: + +```python +mad.send( + """ +if my_condition then + pydbg() +end +""" +) +``` + +### 2.3 Entering the Debugger from `py:send([[...]])` + +The Python execution bridge also exposes `breakpoint` and `pydbg` when using {meth}`MAD.recv_and_exec`. + +```python +mad.send("py:send([[breakpoint()]])") +mad.recv_and_exec() +``` + +This is useful when MAD-NG wants to request a Python-driven debugging pause without embedding the control flow on the Python side ahead of time. + +You can also use the alias: + +```python +mad.send("py:send([[pydbg()]])") +mad.recv_and_exec() +``` + +### 2.4 Scripted Debugger Commands + +For tests, automation, or non-interactive environments, pass a list of debugger commands: + +```python +mad.breakpoint(commands=["h", "c"]) +``` + +This sends: + +1. `h` to print help +2. `c` to continue execution + +This mode is especially useful in unit tests, CI jobs, and examples where a real terminal is not available. + +### 2.5 Continue vs Quit + +The MAD debugger has two very different exit modes: + +- `c` / `continue`: resume execution and keep the MAD process alive +- `q` / `quit`: terminate the MAD process + +If you quit from the debugger, PyMAD-NG closes its side of the pipes cleanly and returns from {meth}`MAD.breakpoint` without raising a traceback. After that, the `MAD` session is finished and should be recreated before sending more commands. + +### 2.6 Practical Example + +```python +from pymadng import MAD + +with MAD() as mad: + mad.send("a = 2") + mad.send("b = 3") + mad.send( + """ +function inspect_sum() + local total = a + b + if total == 5 then + breakpoint() + end + py:send(total) +end +inspect_sum() +""" + ) + print(mad.recv()) +``` + +### 2.7 Terminal Limitations + +The debugger bridge works over the same pipes used for normal PyMAD-NG communication. Because of that: + +- MAD-NG itself is **not** attached to your terminal directly +- line editing is provided on the Python side when a terminal is available +- in non-terminal environments, use `commands=[...]` or a basic line-oriented stream + +If you are in a notebook or another non-TTY environment, prefer scripted commands: + +```python +mad.breakpoint(commands=["where", "c"]) +``` + +--- + +## 3. Inspecting Command History PyMAD-NG keeps track of all **string-based commands** sent to MAD-NG in a history buffer. To review them: @@ -57,9 +191,9 @@ Binary data (like large NumPy arrays) won’t appear in `mad.history()`. Only te --- -## 3. Communication Rules +## 4. Communication Rules -### 3.1 Send Before Receive +### 4.1 Send Before Receive PyMAD-NG uses pipes for **first-in-first-out** communication. If you call: @@ -75,7 +209,7 @@ without telling MAD-NG to `py:send(...)` first, your script will hang. Any mismatch in these calls can lead to deadlocks. -### 3.2 Matching Data Transfers +### 4.2 Matching Data Transfers If you instruct MAD-NG to receive data (`arr = py:recv()`), you must ensure Python **actually sends** that data: @@ -88,9 +222,9 @@ Failing to do so can cause indefinite blocking or partial reads. --- -## 4. Handling Errors +## 5. Handling Errors -### 4.1 Protected Sends +### 5.1 Protected Sends All sends in PyMAD-NG are automatically “protected” by default. If MAD-NG issues an error (`err_`), PyMAD-NG raises a `RuntimeError` on the Python side. @@ -111,39 +245,42 @@ and manually check for failures. --- -## 5. Debugging Subprocess Behavior +## 6. Debugging Subprocess Behaviour -### 5.1 Starting MAD-NG +### 6.1 Starting MAD-NG During initialisation, PyMAD-NG calls: ``` -mad_binary -q -e "MAD.pymad 'py' {_dbg = true} :__ini(fd)" +mad_binary -q -e "" ``` - If `mad_binary` is missing or not executable, you’ll get `FileNotFoundError`. - If it fails to run, an `OSError` is raised. +- The startup chunk now also exports MAD-side debugger aliases (`python_breakpoint`, `pydbg`, and `breakpoint`) before entering the pipe loop. -### 5.2 Checking Streams +### 6.2 Checking Streams - **stdout**: By default prints to Python’s standard output unless you pass `stdout=...`. -- **stderr**: Remains attached to Python’s own stderr, unless you specify `redirect_sterr=True`. +- **stderr**: Remains attached to Python’s own stderr, unless you specify `redirect_stderr=True`. --- -## 6. Common Pitfalls & Solutions +## 7. Common Pitfalls & Solutions | Issue | Possible Cause | Recommended Fix | |------------------------------------------|------------------------------------------------------------------|-----------------------------------------------------------| | **Hang / Deadlock** | Called `mad.recv()` without `mad.send(...)`, or vice versa | Always pair `send()` → `recv()`. Use `debug=True` to see if MAD is expecting data. | | **BrokenPipeError** | MAD-NG crashed or closed unexpectedly | Re-initialise `MAD()`. Check logs for the underlying error. | +| **Debugger `q` ended the session** | `q` in `MAD.dbg()` terminates the MAD subprocess | Create a new `MAD()` instance if you want to continue working. | +| **Arrow keys or history unavailable** | Running without a real terminal | Use a standard terminal or scripted `commands=[...]`. | | **“Unsupported data type”** error | Attempted to `send()` an object that PyMAD-NG can’t serialise | Limit data to `str`, `int`, `float`, `bool`, `list`, or `np.ndarray`. | | **AttributeError / KeyError** accessing a field | Tried to read a reference property without evaluating it first | Call `.eval()` if you need the actual value. | | **Exceeding `_last[]`** references | Too many temp variables stored in `_last[]` | Manually name them in MAD, or increase `num_temp_vars`. | --- -## 7. Cleaning Up +## 8. Cleaning Up If you’re done using MAD-NG, **close** the session: @@ -162,9 +299,12 @@ with MAD(debug=True) as mad: --- -## 8. Summary +## 9. Summary - **Enable** `debug=True` to see more logs. +- **Use** `mad.breakpoint()` to drop into `MAD.dbg()` from Python. +- **Call** `breakpoint()` or `pydbg()` inside MAD strings when you want MAD-side code to trigger debugging. +- **Trigger** `breakpoint()` or `pydbg()` inside `py:send([[...]])` blocks when using `mad.recv_and_exec()`. - **Check** `mad.history()` to identify incorrect or unexpected commands. - **Balance** each `mad.send()` with a `mad.recv()` to avoid deadlocks. - **Catch** `RuntimeError` to handle failures gracefully. @@ -176,4 +316,3 @@ If you still have trouble: - Open an issue on GitHub if you suspect a bug in the code. Happy debugging! - diff --git a/docs/source/ex-lhc-couplingLocal.rst b/docs/source/ex-lhc-couplingLocal.rst index ed8d97c..7639ae9 100644 --- a/docs/source/ex-lhc-couplingLocal.rst +++ b/docs/source/ex-lhc-couplingLocal.rst @@ -7,10 +7,10 @@ LHC Example The file :ref:`ex-lhc-couplingLocal/ex-lhc-couplingLocal.py ` contains an example of loading the required files to use and run the LHC, while including a method to receive and plot intermediate results of a match. -Loading the LHC +Loading the LHC --------------- -The following lines loads the required variables and files for the example. ``assertf`` is not required, but it is used to check if the loading of the LHC is successful. The two string inputs to ``mad.MADX.load`` is also not required, but it is used to specify the final destination of the translated files from MAD-X to MAD-NG. Also, line 5 is not required as this is just to prevent MAD-NG from reporting warnings of unset variables from loading the LHC. Once again, in this extract, the important point to note is that to input strings into functions in MAD-NG, they must be enclosed in quotes. This is because any string input from the side of python is evaluated by MAD-NG, therefore the input can be a variable or expressions. +The following lines load the required variables and files for the example. ``assertf`` is not required, but it is used to check whether loading the LHC is successful. The two string inputs to ``mad.MADX.load`` are also not required, but they are used to specify the final destination of the translated files from MAD-X to MAD-NG. Also, line 5 is not required as this is just to prevent MAD-NG from reporting warnings about unset variables while loading the LHC. Once again, in this extract, the important point to note is that to input strings into functions in MAD-NG, they must be enclosed in quotes. This is because any string input from the Python side is evaluated by MAD-NG, so the input can be a variable or an expression. To grab variables from the MAD-X environment, we use ``mad.load("MADX", ...)``. @@ -19,17 +19,17 @@ To grab variables from the MAD-X environment, we use ``mad.load("MADX", ...)``. :linenos: Receiving intermediate results ------------------------------ +------------------------------ -The most complicated part of the example includes the following set of lines. +The most complicated part of the example includes the following set of lines. From lines 4 - 8 below, we define a function that will be invoked during the optimisation process at each iteration. Within this function, we perform a twiss for the match function to use, while also sending some information on the twiss to python, on line 6. -From lines 10 - 23, we run a match, with a **reference** to the match result returned to the variable ``match_rtrn``. Line 24 is a very important line, as this is something you place in the pipe to MAD-NG for MAD-NG to execute once the match is done. Lines 23-25 receive the first result returned during the match, so that we can start plotting the results. +From lines 10 - 23, we run a match, with a **reference** to the match result returned to the variable ``match_rtrn``. Line 24 is a very important line, as this is something you place in the pipe to MAD-NG for MAD-NG to execute once the match is done. Lines 23 - 25 receive the first result returned during the match, so that we can start plotting the results. -The plotting occurs between lines 29 - 38, wtih the while loop continuing until twiss result is ``None``, which occurs when the match is done, as requested on line 24. +The plotting occurs between lines 29 - 38, with the while loop continuing until the twiss result is ``None``, which occurs when the match is done, as requested on line 24. -Finally, on lines 40 and 41, we retrieve the results of the match from the variable ``match_rtrn``. Since ``match_rtrn`` is a *temporary variable*, there is a limit to how many of these that can be stored (see :doc:`/advanced_features` for more information on these), we delete the reference in python to clear the temporary variable so that is is available for future use. +Finally, on lines 40 and 41, we retrieve the results of the match from the variable ``match_rtrn``. Since ``match_rtrn`` is a *temporary variable*, there is a limit to how many of these can be stored (see :doc:`/advanced_features` for more information), we delete the reference in Python to clear the temporary variable so that it is available for future use. .. important:: As MAD-NG is running in the background, the variable ``match_rtrn`` contains *no* information and instead must be queried for the results. During the query, python will then have to wait for MAD-NG to finish the match, and then return the results. On the other hand, if we do not query for the results, the match will continue to run in the background, we can do other things in python, and then query for the results later. @@ -51,12 +51,12 @@ Finally, on lines 40 and 41, we retrieve the results of the match from the varia .. :lines: 44-47 -.. The ``reg_expr`` fucntion is recursive so, after calling the function, to ensure Python stays in sync, Python is required to ask MAD-NG for the string ``"done"``. +.. The ``reg_expr`` function is recursive so, after calling the function, to ensure Python stays in sync, Python is required to ask MAD-NG for the string ``"done"``. .. .. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py .. :lines: 50-53 -.. Next, we have the first of two methods to evaluate every deferred expression in the LHC and receive them, where Python performs a loop through the number of deferred expressions and has to ask MAD-NG everytime to receive the result. +.. Next, we have the first of two methods to evaluate every deferred expression in the LHC and receive them, where Python performs a loop through the number of deferred expressions and has to ask MAD-NG every time to receive the result. .. .. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py .. :lines: 60-64 @@ -66,7 +66,7 @@ Finally, on lines 40 and 41, we retrieve the results of the match from the varia .. .. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py .. :lines: 67-71 -.. The two methods above do not store the data, so the next bit of code is identical to above, but uses list comprehension to store the data into a list automatically, storing the lists into variables ``exprList1`` and ``exprList2``. The main point of seperating the methods above and below was to identify if storing the variables into a list was a bottleneck. +.. The two methods above do not store the data, so the next bit of code is identical to above, but uses list comprehension to store the data into a list automatically, storing the lists into variables ``exprList1`` and ``exprList2``. The main point of separating the methods above and below was to identify whether storing the variables into a list was a bottleneck. .. .. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py .. :lines: 74-77, 79-83 @@ -101,4 +101,4 @@ Finally, on lines 40 and 41, we retrieve the results of the match from the varia .. .. code-block:: console .. time to retrieve every element name in lhcb1 sequence 0.024236202239990234 sec -.. time to retrieve every element name in lhcb2 sequence 0.0245511531829834 sec \ No newline at end of file +.. time to retrieve every element name in lhcb2 sequence 0.0245511531829834 sec diff --git a/docs/source/function_reference.md b/docs/source/function_reference.md index 9131930..a5f7ab4 100644 --- a/docs/source/function_reference.md +++ b/docs/source/function_reference.md @@ -117,9 +117,49 @@ The first argument is the module path. Optional additional arguments allow impor | {func}`MAD.recv` | Receive results or values from MAD-NG | | {func}`MAD.eval` | Evaluate an expression and return the result | | {func}`MAD.recv_and_exec`| Execute Python code sent from MAD-NG | +| {func}`MAD.breakpoint` | Enter the MAD-NG debugger from Python | +| {func}`MAD.pydbg` | Alias for {func}`MAD.breakpoint` | These tools are essential for controlling the MAD subprocess directly. +### Interactive Debugger Bridge + +PyMAD-NG exposes MAD-NG's `MAD.dbg()` debugger through the high-level interface. + +#### From Python + +```python +mad.breakpoint() +mad.pydbg() +``` + +#### From MAD strings + +PyMAD-NG installs MAD-side helper functions at startup: + +```python +mad.send("breakpoint()") +mad.send("pydbg()") +mad.send("python_breakpoint()") +``` + +#### From `recv_and_exec()` + +When executing Python code sent from MAD-NG, the execution environment automatically includes `breakpoint` and `pydbg`: + +```python +mad.send("py:send([[breakpoint()]])") +mad.recv_and_exec() +``` + +For automated debugging sessions, pass scripted commands: + +```python +mad.breakpoint(commands=["where", "c"]) +``` + +If the debugger command is `q`, the MAD subprocess terminates and the current `MAD` session is finished. + --- ## Reference Objects and Evaluation @@ -198,4 +238,4 @@ mad.tbl.write("'output.tfs'", mad.quote_strings(columns)) PyMAD-NG makes many MAD-NG modules and tools accessible from Python in an intuitive way. This reference highlights the most useful global objects, simulation modules, utility functions, and internal submodules available for scripting or extension work. -For additional control or automation, refer to the API reference or the examples directory. \ No newline at end of file +For additional control or automation, refer to the API reference or the examples directory. diff --git a/docs/source/installation.md b/docs/source/installation.md index d9ef8ad..93523a8 100644 --- a/docs/source/installation.md +++ b/docs/source/installation.md @@ -135,7 +135,7 @@ except RuntimeError as e: Or it is possible to ignore the error and continue execution of the python script: ```python -mad = MAD(raise_on_error=False) +mad = MAD(raise_on_madng_error=False) mad.send("a = 1/'0'") assert mad.send("py:send(1)").recv() == 1 ``` @@ -157,6 +157,25 @@ mad.send("py:send('print(\'debug line\')')") mad.recv_and_exec() ``` +The execution context also exposes the debugger bridge, so MAD can ask Python to open the MAD debugger: + +```python +mad.send("py:send([[breakpoint()]])") +mad.recv_and_exec() +``` + +You can also stop from Python directly: + +```python +mad.breakpoint() +``` + +For scripted or non-interactive sessions: + +```python +mad.breakpoint(commands=["h", "c"]) +``` + ### 6. Ensure Communication Order MAD-NG uses FIFO communication. Always follow the rule: @@ -174,4 +193,3 @@ Now that PyMAD-NG is installed, you can move on to: - **[Quick Start Guide →](quickstartguide.md)** *(Learn the basics and run your first PyMAD-NG script)* - **[API Reference →](reference.rst)** *(Explore the functions and classes available in PyMAD-NG)* - diff --git a/docs/source/quickstartguide.md b/docs/source/quickstartguide.md index 09fe983..2d6677e 100644 --- a/docs/source/quickstartguide.md +++ b/docs/source/quickstartguide.md @@ -20,6 +20,18 @@ mad = MAD() This automatically launches a MAD-NG process and connects it to Python. +```{important} +`MAD()` launches a real MAD-NG subprocess immediately. For normal scripts, prefer `with MAD() as mad:` so the process is always cleaned up. +``` + +### Mental Model + +Four rules explain most PyMAD-NG behaviour: + +1. `mad.send("...")` sends code or data into the running MAD-NG process. +2. `mad.recv()` only succeeds after MAD-NG has been told to `py:send(...)` something back. +3. High-level objects like `mad.seq`, `mad.twiss(...)`, or `mad.math.sin(1)` are usually references into MAD-NG, not copied Python values. + --- ## Step 2: Load a Sequence @@ -58,6 +70,28 @@ mad.send("seq.beam = beam {}") --- +## Step 3.5: References vs Values + +PyMAD-NG keeps most large or structured objects inside MAD-NG and gives Python a reference to them. + +```python +tbl_ref = mad.twiss(sequence=mad.seq) +print(type(tbl_ref)) +``` + +That is useful for performance and for symbolic chaining, but it also means you should explicitly ask for concrete Python values when needed: + +```python +value = mad.math.sin(1).eval() +df = mad.tbl.to_df() +``` + +```{note} +If an object feels more like a handle than a Python value, that is usually correct. Use `.eval()` or a conversion method such as `.to_df()` when you want data materialised on the Python side. +``` + +--- + ## Step 4: Run a Twiss Calculation ### High-Level: @@ -76,6 +110,25 @@ tbl = mad.recv() --- +## Step 4.5: Assigning MAD-NG Variables Correctly + +To create or overwrite a variable in MAD-NG, use square brackets: + +```python +mad["energy"] = 6500 +mad["seq"] = mad.MADX.seq +``` + +Do not use: + +```python +mad.energy = 6500 +``` + +That only changes the Python wrapper object and does not update the running MAD-NG process. + +--- + ## Step 5: Analyse the Results Convert the resulting MAD table to a Pandas DataFrame: @@ -111,7 +164,18 @@ plt.show() print(mad.recv()) ``` -- Always match `send()` and `recv()` properly to avoid blocking communication. +- If MAD-NG is expecting Python data, send it immediately: + ```python + mad.send("arr = py:recv()") + mad.send(my_array) + ``` + +- You can enter the MAD debugger from Python: + ```python + mad.breakpoint() + ``` + +- Quitting the debugger with `q` terminates the current MAD session. After that, create a new `MAD()` instance before continuing. ## What Next? @@ -119,4 +183,3 @@ Now that you’ve completed your first PyMAD-NG workflow, explore: - **[MAD-NG Documentation](https://madx.web.cern.ch/releases/madng/html/)** for details on MAD-NG features - **[API Reference](reference.rst)** for full documentation of the {class}`pymadng.MAD` class - **[Examples](examples.rst)** to see real-world scripts. With all the scripts also in the [`examples` folder of the PyMAD-NG repository](https://github.com/MethodicalAcceleratorDesign/MAD-NG.py/tree/main/examples). - diff --git a/pyproject.toml b/pyproject.toml index d15aefe..dd59d45 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,7 +38,9 @@ where = ["src"] version = {attr = "pymadng.__version__"} [project.optional-dependencies] +dev = ["pre-commit>=4,<5", "ruff>=0.13,<0.14"] tfs = ["tfs-pandas>3.0.0"] +test = ["pytest>=8,<9", "pytest-cov>=6,<7"] # ----- Dev Tools Configuration ----- # @@ -91,4 +93,16 @@ extend-select = [ ] # Allow fix for all enabled rules (when `--fix`) is provided. fixable = ["ALL"] -unfixable = [] \ No newline at end of file +unfixable = [] + +[tool.pytest.ini_options] +testpaths = ["tests"] +python_files = ["test_*.py"] +addopts = "-ra" + +[tool.coverage.run] +relative_files = true +source = ["src/pymadng"] + +[tool.coverage.report] +omit = ["tests/*"] diff --git a/src/pymadng/__init__.py b/src/pymadng/__init__.py index db957dd..949a982 100644 --- a/src/pymadng/__init__.py +++ b/src/pymadng/__init__.py @@ -1,7 +1,7 @@ from .madp_object import MAD __title__ = "pymadng" -__version__ = "0.9.0" +__version__ = "0.10.0" __summary__ = "Python interface to MAD-NG running as subprocess" __uri__ = "https://github.com/MethodicalAcceleratorDesign/MAD-NG.py" diff --git a/src/pymadng/madp_object.py b/src/pymadng/madp_object.py index 0432ea3..1ac2ad5 100644 --- a/src/pymadng/madp_object.py +++ b/src/pymadng/madp_object.py @@ -172,7 +172,7 @@ def receive(self, varname: str | None = None) -> Any: """ return self.__process.recv(varname) - def recv_and_exec(self, context: dict = {}) -> dict: + def recv_and_exec(self, context: dict | None = None) -> dict: """ Receive a string from MAD-NG and execute it. @@ -184,9 +184,37 @@ def recv_and_exec(self, context: dict = {}) -> dict: Returns: dict: The updated execution environment. """ + if context is None: + context = {} context["mad"] = self + context.setdefault("breakpoint", self.breakpoint) + context.setdefault("pydbg", self.breakpoint) return self.__process.recv_and_exec(context) + def breakpoint( + self, commands: Iterable[str] | None = None, input_stream: TextIO | None = None + ) -> MAD: + """ + Enter the MAD-NG debugger through the active subprocess. + + Args: + commands (Iterable[str] | None, optional): Scripted debugger commands. When + omitted, commands are read interactively from the terminal. + input_stream (TextIO | None, optional): Alternate stream to read debugger + commands from in interactive mode. + + Returns: + MAD: Self for method chaining once the debugger resumes. + """ + self.__process.enter_debugger(commands=commands, input_stream=input_stream) + return self + + def pydbg( + self, commands: Iterable[str] | None = None, input_stream: TextIO | None = None + ) -> MAD: + """Alias for :meth:`breakpoint`.""" + return self.breakpoint(commands=commands, input_stream=input_stream) + # --------------------------------Sending data to subprocess------------------------------------# def send(self, data: Any) -> MAD: """ @@ -331,9 +359,7 @@ def loadfile(self, path: str | Path, *varnames: str): """ path: Path = Path(path).resolve() if varnames == (): - self.__process.send( - f"assert(loadfile('{path}', nil, {self.py_name}._env))()" - ) + self.__process.send(f"assert(loadfile('{path}', nil, {self.py_name}._env))()") else: # The parent/stem is necessary, otherwise the file will not be found # This is thanks to the way the require function works in MAD-NG (how it searches for files) diff --git a/src/pymadng/madp_pymad.py b/src/pymadng/madp_pymad.py index ca5df4c..9fed095 100644 --- a/src/pymadng/madp_pymad.py +++ b/src/pymadng/madp_pymad.py @@ -15,12 +15,18 @@ import numpy as np if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Iterable from numpy.typing import DTypeLike # TODO: look at cpymad for the suppression of the error messages at exit - copy? (jgray 2024) +DEBUGGER_ENTER_MESSAGE = "__pymadng_debug_enter__" +DEBUGGER_EXIT_MESSAGE = "__pymadng_debug_exit__" +DEBUGGER_RESUME_COMMANDS = {"c", "cont", "continue"} +DEBUGGER_TERMINATE_COMMANDS = {"q", "quit"} +DEBUGGER_PROMPT = "\001\033[34m\002dbg>\001\033[0m\002 " + def is_private(varname): """Check if the variable name is considered private. @@ -80,11 +86,22 @@ def __init__( # Redirect stderr to stdout, if specified stderr = stdout if redirect_stderr else sys.stderr.fileno() - # Create a chunk of code to start the process lua_debug_flag = "true" if debug else "false" - startup_chunk = ( - f"MAD.pymad '{py_name}' {{_dbg = {lua_debug_flag}}} :__ini({mad_write})" - ) + startup_chunk = f""" +local __pymad = MAD.pymad '{py_name}' {{_dbg = {lua_debug_flag}}} +rawset(_G, '{py_name}', __pymad) +local function __python_breakpoint() + __pymad:send('{DEBUGGER_ENTER_MESSAGE}') + io.stdout:setvbuf('no') + MAD.dbg() + io.stdout:setvbuf('line') + __pymad:send('{DEBUGGER_EXIT_MESSAGE}') +end +rawset(_G, 'python_breakpoint', __python_breakpoint) +rawset(_G, 'pydbg', __python_breakpoint) +rawset(_G, 'breakpoint', __python_breakpoint) +__pymad:__ini({mad_write}) +""" if threading.current_thread() is threading.main_thread(): self._setup_signal_handler() @@ -114,6 +131,7 @@ def __init__( self.mad_read_stream = os.fdopen(self.mad_output_pipe, "rb") self.history = "" # Begin the recording of the history self.debug = debug # Record debug mode status + self._debugger_active = False # stdout should be line buffered by default, but for jupyter notebook, # stdout is redirected and not line buffered by default @@ -130,9 +148,7 @@ def __init__( if not startup_status_checker[0] or mad_rtrn != "started": self.close() if mad_rtrn == "started": - raise OSError( - f"Could not establish communication with {mad_path} process" - ) + raise OSError(f"Could not establish communication with {mad_path} process") raise OSError(f"Could not start {mad_path} process, received: {mad_rtrn}") # Set the error handler to be on by default @@ -150,6 +166,155 @@ def delete_process(sig, frame): signal.signal(signal.SIGINT, delete_process) + def _normalise_debugger_command(self, command: str) -> bytes: + if not isinstance(command, str): + raise TypeError("Debugger commands must be strings") + if not command.endswith("\n"): + command += "\n" + return command.encode("utf-8") + + def _send_debugger_command(self, command: str) -> None: + self.mad_input_stream.write(self._normalise_debugger_command(command)) + + def _read_debugger_state(self, timeout: float = 0.1) -> str: + if self.process.poll() is not None: + return "terminated" + + ready, _, _ = select.select([self.mad_read_stream], [], [], timeout) + if not ready: + return "active" + + try: + message = self.recv() + except Exception: + if self.process.poll() is not None: + return "terminated" + raise + + if isinstance(message, BrokenPipeError): + return "terminated" + + if message == DEBUGGER_EXIT_MESSAGE: + return "resumed" + + raise RuntimeError(f"Unexpected message received while debugging: {message!r}") + + def enter_debugger( + self, commands: Iterable[str] | None = None, input_stream: TextIO | None = None + ) -> None: + if self._debugger_active: + raise RuntimeError("MAD debugger is already active") + + scripted_commands = None + if commands is not None: + scripted_commands = list(commands) + if not scripted_commands: + raise ValueError("Debugger command list must not be empty") + + self._debugger_active = True + try: + self.send("python_breakpoint()") + entry_message = self.recv() + if entry_message != DEBUGGER_ENTER_MESSAGE: + raise RuntimeError( + f"Unexpected message received while entering debugger: {entry_message!r}" + ) + + if scripted_commands is not None: + self._run_scripted_debugger(scripted_commands) + return + + self._run_interactive_debugger(input_stream) + finally: + self._debugger_active = False + + def _run_scripted_debugger(self, commands: list[str]) -> None: + final_command = commands[-1].strip().split(maxsplit=1)[0].lower() + if ( + final_command not in DEBUGGER_RESUME_COMMANDS + and final_command not in DEBUGGER_TERMINATE_COMMANDS + ): + raise ValueError( + "Scripted debugger commands must end with continue ('c') or quit ('q')" + ) + + for command in commands: + self._send_debugger_command(command) + state = self._read_debugger_state() + if state == "resumed": + return + if state == "terminated": + self.close() + return + + state = self._read_debugger_state(timeout=1.0) + if state == "resumed": + return + if state == "terminated": + self.close() + return + raise RuntimeError("MAD debugger did not resume after scripted commands") + + def _run_interactive_debugger(self, input_stream: TextIO | None) -> None: + if input_stream is None and self._stdin_is_tty(): + self._run_readline_debugger() + return + + if input_stream is None: + try: + with Path("/dev/tty").open() as tty_stream: + self._run_interactive_debugger(tty_stream) + return + except OSError: + input_stream = sys.stdin + + while True: + command = input_stream.readline() + if command == "": + if self.process.poll() is not None: + raise RuntimeError("MAD debugger terminated the subprocess") + raise EOFError("Reached EOF while the MAD debugger was active") + + self._send_debugger_command(command) + state = self._read_debugger_state() + if state == "resumed": + return + if state == "terminated": + self.close() + return + + def _stdin_is_tty(self) -> bool: + try: + return os.isatty(sys.stdin.fileno()) + except (AttributeError, OSError, ValueError): + return False + + def _run_readline_debugger(self) -> None: + while True: + try: + self._render_debugger_prompt() + command = input(DEBUGGER_PROMPT) + except EOFError as exc: + if self.process.poll() is not None: + raise RuntimeError("MAD debugger terminated the subprocess") from exc + raise EOFError("Reached EOF while the MAD debugger was active") from exc + + self._send_debugger_command(command) + state = self._read_debugger_state() + if state == "resumed": + return + if state == "terminated": + self.close() + return + + def _render_debugger_prompt(self) -> None: + try: + stdout = sys.stdout + stdout.write("\r\x1b[2K") + stdout.flush() + except (AttributeError, OSError, ValueError): + pass + def send_range(self, start: float, stop: float, size: int) -> None: """Send a linear range (numpy array) to MAD-NG. @@ -206,13 +371,9 @@ def protected_send(self, string: str) -> MadProcess: if self.raise_on_madng_error: # If the user has specified that they want to raise an error always, skip the error handling on and off return self.send(string) - return self.send( - f"{self.py_name}:__err(true); {string}; {self.py_name}:__err(false);" - ) + return self.send(f"{self.py_name}:__err(true); {string}; {self.py_name}:__err(false);") - def protected_variable_retrieval( - self, name: str, shallow_copy: bool = False - ) -> Any: + def protected_variable_retrieval(self, name: str, shallow_copy: bool = False) -> Any: """Safely retrieve a variable from MAD-NG. Enables temporary error handling while retrieving a variable. @@ -257,7 +418,7 @@ def recv(self, varname: str | None = None) -> Any: self.varname = varname # For mad reference return type_fun[typ]["recv"](self) # type: ignore - def recv_and_exec(self, env: dict = {}) -> dict: + def recv_and_exec(self, env: dict | None = None) -> dict: """Receive a command string from MAD-NG and execute it. The execution context includes numpy as np and the mad process instance. @@ -267,12 +428,18 @@ def recv_and_exec(self, env: dict = {}) -> dict: Returns: dict: The updated environment dictionary after executing the received command. """ + if env is None: + env = {} + # Check if user has already defined mad (madp_object will have mad defined), otherwise define it try: env["mad"] except KeyError: env["mad"] = self + env.setdefault("breakpoint", self.enter_debugger) + env.setdefault("pydbg", self.enter_debugger) + exec(compile(self.recv(), "ffrom_mad", "exec"), self.python_exec_context, env) return env @@ -309,9 +476,7 @@ def recv_vars(self, *names, shallow_copy: bool = False) -> Any: raise ValueError("Cannot retrieve private variables from MAD-NG") if len(names) == 1: return self.protected_variable_retrieval(names[0], shallow_copy) - return tuple( - self.protected_variable_retrieval(name, shallow_copy) for name in names - ) + return tuple(self.protected_variable_retrieval(name, shallow_copy) for name in names) # -------------------------------------------------------------------------- # @@ -321,15 +486,18 @@ def close(self) -> None: Closes all communication pipes and waits for the subprocess to finish. """ if self.process.poll() is None: # If process is still running - self.send(f"{self.py_name}:__fin()") # Tell the mad side to finish - open_pipe = select.select([self.mad_read_stream], [], [], 5) # Shorter timeout - if open_pipe[0]: - # Wait for the mad side to finish - close_msg = self.recv("closing") - if close_msg != "": - logging.warning( - f"Unexpected message received: {close_msg}, MAD-NG may not have completed properly" - ) + try: + self.send(f"{self.py_name}:__fin()") # Tell the mad side to finish + open_pipe = select.select([self.mad_read_stream], [], [], 5) # Shorter timeout + if open_pipe[0]: + # Wait for the mad side to finish + close_msg = self.recv("closing") + if close_msg != "": + logging.warning( + f"Unexpected message received: {close_msg}, MAD-NG may not have completed properly" + ) + except BrokenPipeError: + pass # Always try terminate first self.process.terminate() @@ -702,9 +870,7 @@ def recv_generic_matrix(self: MadProcess, dtype: np.dtype) -> np.ndarray: np.ndarray: The received matrix as a reshaped numpy array. """ shape = read_data_stream(self, 8, np.int32) - return read_data_stream(self, shape[0] * shape[1] * dtype.itemsize, dtype).reshape( - shape - ) + return read_data_stream(self, shape[0] * shape[1] * dtype.itemsize, dtype).reshape(shape) def recv_matrix(self: MadProcess) -> np.ndarray: @@ -780,9 +946,7 @@ def send_generic_tpsa( assert len(monos) == len(coefficients), ( "The number of monomials must be equal to the number of coefficients" ) - assert monos.dtype == np.uint8, ( - "The monomials must be of type 8-bit unsigned integer " - ) + assert monos.dtype == np.uint8, "The monomials must be of type 8-bit unsigned integer " write_serial_data(self, "ii", len(monos), len(monos[0])) for mono in monos: self.mad_input_stream.write(mono.tobytes()) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/inputs/example.log b/tests/inputs/example.log index 3a54aab..98e9ffe 100644 --- a/tests/inputs/example.log +++ b/tests/inputs/example.log @@ -27,7 +27,7 @@ match = MAD.match ***pymad.recv: [py:__err(true):send(MAD['env']['version'], false):__err(false)] 62 bytes ***pymad.send: [str_] 4 bytes ***pymad.send: binary data 4 bytes -***pymad.send: [1.1.11] 6 bytes +***pymad.send: [1.1.10] 6 bytes ***pymad.recv: binary data 4 bytes ***pymad.recv: [ function __mklast__ (a, b, ...) diff --git a/tests/test_communication.py b/tests/test_communication.py index 14d7d26..a33aabd 100644 --- a/tests/test_communication.py +++ b/tests/test_communication.py @@ -1,93 +1,98 @@ +from __future__ import annotations + import threading -import unittest from unittest.mock import patch +import pytest + from pymadng import MAD -class TestExecution(unittest.TestCase): - def test_recv_and_exec(self): - with MAD() as mad: - mad.send( - """py:send([==[mad.send('''py:send([=[mad.send("py:send([[a = 100/2]])")]=])''')]==])""" - ) - mad.recv_and_exec() - mad.recv_and_exec() - a = mad.recv_and_exec()["a"] - self.assertEqual(a, 50) - - def test_err(self): - with MAD(stdout="/dev/null", redirect_stderr=True) as mad: - mad.send("py:__err(true)") - mad.send("1+1") # Load error - self.assertRaises(RuntimeError, mad.recv) - mad.send("py:__err(true)") - mad.send("print(nil/2)") # Runtime error - self.assertRaises(RuntimeError, mad.recv) - - -class TestStrings(unittest.TestCase): - def test_recv(self): - with MAD() as mad: - mad.send("py:send('hi')") - mad.send("""py:send([[Multiline string should work +def test_recv_and_exec(): + with MAD() as mad: + mad.send( + """py:send([==[mad.send('''py:send([=[mad.send("py:send([[a = 100/2]])")]=])''')]==])""" + ) + mad.recv_and_exec() + mad.recv_and_exec() + assert mad.recv_and_exec()["a"] == 50 + + +def test_err(): + with MAD(stdout="/dev/null", redirect_stderr=True) as mad: + mad.send("py:__err(true)") + mad.send("1+1") + with pytest.raises(RuntimeError): + mad.recv() + + mad.send("py:__err(true)") + mad.send("print(nil/2)") + with pytest.raises(RuntimeError): + mad.recv() + + +def test_recv(): + with MAD() as mad: + mad.send("py:send('hi')") + mad.send("""py:send([[Multiline string should work Like So.]])""") - self.assertEqual(mad.recv(), "hi") - self.assertEqual(mad.receive(), "Multiline string should work\n\nLike So.") - - def test_send(self): - with MAD() as mad: - init_str = "asdfghjkl;" - mad.send("str = py:recv(); py:send(str .. str)") - mad.send(init_str) - self.assertEqual(mad.recv(), init_str * 2) - mad.send("str2 = py:recv(); py:send(str2 .. str2)") - init_str = """Py Multiline string should work - -Like So.]])""" - mad.send(init_str) - self.assertEqual(mad.recv(), init_str * 2) - - def test_protected_send(self): - with MAD(stdout="/dev/null", redirect_stderr=True, raise_on_madng_error=False) as mad: - mad.send("py:send('hello world'); a = nil/2") - self.assertEqual(mad.recv(), "hello world") # python should not crash - mad.send("py:send(1)") - self.assertEqual(mad.recv(), 1) # Check that the error did not affect the pipe - - mad.protected_send("a = nil/2") - self.assertRaises(RuntimeError, mad.recv) # python should receive an error - - mad.psend("a = nil/2") - self.assertRaises(RuntimeError, mad.recv) - - -class TestOutput(unittest.TestCase): - def test_print(self): - with MAD() as mad: - mad.send("py:send('hello world')") - self.assertEqual(mad.recv(), "hello world") # Check printing does not affect pipe - - -class TestSignalHandler(unittest.TestCase): - @patch("pymadng.madp_pymad.MadProcess._setup_signal_handler") - def test_signal_handler_called_in_main_thread(self, mock_setup_signal_handler): - with patch("pathlib.Path.exists", return_value=True): - MAD() - mock_setup_signal_handler.assert_called_once() + assert mad.recv() == "hi" + assert mad.receive() == "Multiline string should work\n\nLike So." + + +@pytest.mark.parametrize( + "value", + [ + "asdfghjkl;", + """Py Multiline string should work - @patch("pymadng.madp_pymad.MadProcess._setup_signal_handler") - def test_signal_handler_not_called_in_non_main_thread(self, mock_setup_signal_handler): - def create_mad_proc(): - with patch("pathlib.Path.exists", return_value=True): - MAD() +Like So.]])""", + ], +) +def test_send(value: str): + with MAD() as mad: + mad.send("str = py:recv(); py:send(str .. str)") + mad.send(value) + assert mad.recv() == value * 2 - thread = threading.Thread(target=create_mad_proc) - thread.start() - thread.join() - mock_setup_signal_handler.assert_not_called() +def test_protected_send(): + with MAD(stdout="/dev/null", redirect_stderr=True, raise_on_madng_error=False) as mad: + mad.send("py:send('hello world'); a = nil/2") + assert mad.recv() == "hello world" + mad.send("py:send(1)") + assert mad.recv() == 1 + + mad.protected_send("a = nil/2") + with pytest.raises(RuntimeError): + mad.recv() + + mad.psend("a = nil/2") + with pytest.raises(RuntimeError): + mad.recv() + + +def test_print(): + with MAD() as mad: + mad.send("py:send('hello world')") + assert mad.recv() == "hello world" + + +@patch("pymadng.madp_pymad.MadProcess._setup_signal_handler") +def test_signal_handler_called_in_main_thread(mock_setup_signal_handler): + with patch("pathlib.Path.exists", return_value=True): + MAD() + mock_setup_signal_handler.assert_called_once() + + +@patch("pymadng.madp_pymad.MadProcess._setup_signal_handler") +def test_signal_handler_not_called_in_non_main_thread(mock_setup_signal_handler): + def create_mad_proc(): + with patch("pathlib.Path.exists", return_value=True): + MAD() -if __name__ == "__main__": - unittest.main() + thread = threading.Thread(target=create_mad_proc) + thread.start() + thread.join() + mock_setup_signal_handler.assert_not_called() diff --git a/tests/test_debug.py b/tests/test_debug.py index d50ab2f..6725dc4 100644 --- a/tests/test_debug.py +++ b/tests/test_debug.py @@ -1,73 +1,363 @@ +from __future__ import annotations + +import io +import re import time -import unittest from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest from pymadng import MAD +from pymadng.madp_pymad import MadProcess + +INPUTS_FOLDER = Path(__file__).parent / "inputs" +ANSI_ESCAPE = re.compile(r"\x1b\[[0-9;]*m") + + +def strip_ansi(text: str) -> str: + return ANSI_ESCAPE.sub("", text) + + +def make_test_process() -> MadProcess: + process = object.__new__(MadProcess) + process.py_name = "py" + process.process = MagicMock() + process.process.poll.return_value = None + process.mad_read_stream = MagicMock() + process.mad_input_stream = MagicMock() + process.python_exec_context = {} + process._debugger_active = False + process.close = MagicMock() + process.send = MagicMock() + process.recv = MagicMock() + process.protected_variable_retrieval = MagicMock() + return process + + +def test_logfile(tmp_path): + test_log1 = tmp_path / "test.log" + test_log2 = tmp_path / "test2.log" + + with MAD(stdout=test_log1, debug=True, raise_on_madng_error=False): + pass + time.sleep(0.1) + text = test_log1.read_text() + assert "***pymad.recv: binary data 4 bytes" in text + assert "io.stdout:setvbuf('line')" in text + assert "py:send('started')" in text + assert "started" in text + + with MAD(stdout=test_log2, debug=True) as mad: + mad.send("!This is a line that does nothing") + mad.send("print('hello world')") + + text = test_log2.read_text() + assert "[!This is a line that does nothing]" in text + assert "hello world\n" in text + assert "[print('hello world')]" in text + + with MAD(stdout=test_log2, debug=False) as mad: + mad.send("!This is a line that does nothing") + mad.send("print('hello world')") + + assert test_log2.read_text() == "hello world\n" + + +def test_err(tmp_path): + test_log1 = tmp_path / "test.log" + + with ( + test_log1.open("w") as handle, + MAD(debug=True, stdout=handle, raise_on_madng_error=False) as mad, + ): + mad.psend("a = nil/2") + with pytest.raises(RuntimeError): + mad.recv() + + file_text = test_log1.read_text() + assert "[py:__err(true); a = nil/2; py:__err(false);]" in file_text + assert "***pymad.run:" not in file_text + + with MAD(stdout=test_log1, redirect_stderr=True) as mad: + mad.psend("a = nil/2") + with pytest.raises(RuntimeError): + mad.recv() + + assert test_log1.read_text()[:13] == "***pymad.run:" + + +def test_breakpoint(tmp_path): + test_log = tmp_path / "test_breakpoint.log" + + with MAD(stdout=test_log, redirect_stderr=True, raise_on_madng_error=False) as mad: + for name in ("python_breakpoint", "pydbg", "breakpoint"): + mad.send(f"py:send(type({name}))") + assert mad.recv() == "function" + + mad.breakpoint(commands=["h", "c"]) + mad.send("py:send('alive')") + assert mad.recv() == "alive" + + file_text = strip_ansi(test_log.read_text()) + assert "break via dbg()" in file_text + assert "h(elp)" in file_text + + +def test_breakpoint_from_py_send(tmp_path): + test_log = tmp_path / "test_breakpoint_exec.log" + + with MAD(stdout=test_log, redirect_stderr=True, raise_on_madng_error=False) as mad: + mad.send("py:send([[breakpoint(commands=['h', 'c'])]])") + mad.recv_and_exec() + mad.send("py:send([[pydbg(commands=['c'])]])") + mad.recv_and_exec() + mad.send("py:send('alive')") + assert mad.recv() == "alive" + + file_text = strip_ansi(test_log.read_text()) + assert file_text.count("break via dbg()") >= 2 + + +def test_breakpoint_invalid_scripted_commands(tmp_path): + test_log = tmp_path / "test_breakpoint_invalid.log" + + with MAD(stdout=test_log, redirect_stderr=True, raise_on_madng_error=False) as mad: + with pytest.raises(ValueError, match="must not be empty"): + mad.breakpoint(commands=[]) + + mad.send("py:send('alive')") + assert mad.recv() == "alive" + + +def test_breakpoint_interactive_input_stream_resume(tmp_path): + test_log = tmp_path / "test_breakpoint_input_stream_resume.log" + + with MAD(stdout=test_log, redirect_stderr=True, raise_on_madng_error=False) as mad: + mad.breakpoint(input_stream=io.StringIO("h\nc\n")) + mad.send("py:send('alive')") + assert mad.recv() == "alive" + + file_text = strip_ansi(test_log.read_text()) + assert "break via dbg()" in file_text + assert "h(elp)" in file_text + + +def test_breakpoint_interactive_input_stream_quit(tmp_path): + test_log = tmp_path / "test_breakpoint_input_stream_quit.log" + + with MAD(stdout=test_log, redirect_stderr=True, raise_on_madng_error=False) as mad: + mad.breakpoint(input_stream=io.StringIO("q\n")) + assert mad._MAD__process.process.poll() is not None + assert mad._MAD__process.mad_read_stream.closed + assert mad._MAD__process.mad_input_stream.closed + + assert "break via dbg()" in strip_ansi(test_log.read_text()) + + +def test_breakpoint_reentry_from_python_callback(tmp_path): + test_log = tmp_path / "test_breakpoint_reentry.log" + + with MAD(stdout=test_log, redirect_stderr=True, raise_on_madng_error=False) as mad: + mad.send("py:send([[breakpoint(commands=['c'])]])") + mad.recv_and_exec({"breakpoint": lambda *_args, **_kwargs: mad.breakpoint(commands=["c"])}) + + mad.send("py:send('alive')") + assert mad.recv() == "alive" + + +def test_breakpoint_quit(tmp_path): + test_log = tmp_path / "test_breakpoint_quit.log" + + with MAD(stdout=test_log, redirect_stderr=True, raise_on_madng_error=False) as mad: + mad.breakpoint(commands=["q"]) + assert mad._MAD__process.process.poll() is not None + assert mad._MAD__process.mad_read_stream.closed + assert mad._MAD__process.mad_input_stream.closed + + assert "break via dbg()" in strip_ansi(test_log.read_text()) + + +def test_madprocess_startup_failure_paths(): + common_patches = [ + patch("pymadng.madp_pymad.Path.exists", return_value=True), + patch("pymadng.madp_pymad.os.pipe", side_effect=[(10, 11), (12, 13)]), + patch("pymadng.madp_pymad.os.fdopen", side_effect=[MagicMock(), MagicMock()]), + patch("pymadng.madp_pymad.os.close"), + patch("pymadng.madp_pymad.subprocess.Popen", return_value=MagicMock()), + patch.object(MadProcess, "_setup_signal_handler"), + patch.object(MadProcess, "send"), + patch.object(MadProcess, "close"), + ] + + with ( + common_patches[0], + common_patches[1], + common_patches[2], + common_patches[3], + common_patches[4], + common_patches[5], + common_patches[6], + common_patches[7], + patch("pymadng.madp_pymad.select.select", return_value=([], [], [])), + patch.object(MadProcess, "recv", return_value="started"), + pytest.raises(OSError, match="Could not establish communication"), + ): + MadProcess("/tmp/mad") + + with ( + patch("pymadng.madp_pymad.Path.exists", return_value=True), + patch("pymadng.madp_pymad.os.pipe", side_effect=[(10, 11), (12, 13)]), + patch("pymadng.madp_pymad.os.fdopen", side_effect=[MagicMock(), MagicMock()]), + patch("pymadng.madp_pymad.os.close"), + patch("pymadng.madp_pymad.subprocess.Popen", return_value=MagicMock()), + patch.object(MadProcess, "_setup_signal_handler"), + patch.object(MadProcess, "send"), + patch.object(MadProcess, "close"), + patch("pymadng.madp_pymad.select.select", return_value=([object()], [], [])), + patch.object(MadProcess, "recv", return_value="boom"), + pytest.raises(OSError, match="Could not start"), + ): + MadProcess("/tmp/mad") + + +def test_madprocess_internal_debugger_branches(): + process = make_test_process() + with pytest.raises(TypeError, match="Debugger commands must be strings"): + process._normalise_debugger_command(1) # type: ignore[arg-type] + + process = make_test_process() + process.process.poll.return_value = 0 + assert process._read_debugger_state() == "terminated" + + process = make_test_process() + process.recv.side_effect = RuntimeError("boom") + with ( + patch("pymadng.madp_pymad.select.select", return_value=([object()], [], [])), + pytest.raises(RuntimeError, match="boom"), + ): + process._read_debugger_state() + + process = make_test_process() + process.recv.return_value = "unexpected" + with ( + patch("pymadng.madp_pymad.select.select", return_value=([object()], [], [])), + pytest.raises(RuntimeError, match="Unexpected message"), + ): + process._read_debugger_state() + + process = make_test_process() + process._debugger_active = True + with pytest.raises(RuntimeError, match="already active"): + process.enter_debugger() + + process = make_test_process() + process.recv.return_value = "not-enter" + with pytest.raises(RuntimeError, match="Unexpected message received while entering debugger"): + process.enter_debugger(commands=["c"]) + + process = make_test_process() + with pytest.raises(ValueError, match="must end with continue"): + process._run_scripted_debugger(["h"]) + + process = make_test_process() + process._read_debugger_state = MagicMock(side_effect=["active", "active", "resumed"]) + process._run_scripted_debugger(["h", "c"]) + + process = make_test_process() + process._read_debugger_state = MagicMock(side_effect=["active", "active", "terminated"]) + process._run_scripted_debugger(["h", "q"]) + process.close.assert_called_once() + + process = make_test_process() + process._read_debugger_state = MagicMock(side_effect=["active", "active", "active"]) + with pytest.raises(RuntimeError, match="did not resume"): + process._run_scripted_debugger(["h", "c"]) + + process = make_test_process() + process._stdin_is_tty = MagicMock(return_value=True) + process._run_readline_debugger = MagicMock() + process._run_interactive_debugger(None) + process._run_readline_debugger.assert_called_once() + + process = make_test_process() + process._stdin_is_tty = MagicMock(return_value=False) + process._read_debugger_state = MagicMock(return_value="resumed") + with ( + patch("pymadng.madp_pymad.Path.open", side_effect=OSError), + patch("sys.stdin", io.StringIO("c\n")), + ): + process._run_interactive_debugger(None) + + process = make_test_process() + process.process.poll.return_value = 0 + with pytest.raises(RuntimeError, match="terminated the subprocess"): + process._run_interactive_debugger(io.StringIO("")) + + process = make_test_process() + with pytest.raises(EOFError, match="Reached EOF"): + process._run_interactive_debugger(io.StringIO("")) + + process = make_test_process() + with patch("pymadng.madp_pymad.os.isatty", side_effect=ValueError): + assert process._stdin_is_tty() is False + + process = make_test_process() + process._read_debugger_state = MagicMock(return_value="terminated") + with patch("builtins.input", return_value="q"): + process._run_readline_debugger() + process.close.assert_called_once() + + process = make_test_process() + process.process.poll.return_value = 0 + with ( + patch("builtins.input", side_effect=EOFError), + pytest.raises(RuntimeError, match="terminated the subprocess"), + ): + process._run_readline_debugger() + + process = make_test_process() + with ( + patch("builtins.input", side_effect=EOFError), + pytest.raises(EOFError, match="Reached EOF"), + ): + process._run_readline_debugger() + + with patch("sys.stdout.write", side_effect=OSError): + make_test_process()._render_debugger_prompt() + + +def test_madprocess_internal_recv_and_close_branches(): + process = make_test_process() + process.recv.return_value = "value = 3" + env = process.recv_and_exec() + assert env["value"] == 3 + assert env["mad"] is process + + process = make_test_process() + process.recv.return_value = "value = 4" + env = {"mad": "sentinel"} + result = process.recv_and_exec(env) + assert result["mad"] == "sentinel" + + with pytest.raises(ValueError, match="Cannot retrieve private variables"): + make_test_process().recv_vars("_private") + + process = make_test_process() + process.recv.return_value = "unexpected" + process.stdout_file = MagicMock() + process.mad_read_stream.closed = False + process.mad_input_stream.closed = False + with ( + patch("pymadng.madp_pymad.select.select", return_value=([object()], [], [])), + patch("pymadng.madp_pymad.logging.warning") as warning, + ): + MadProcess.close(process) + warning.assert_called_once() -inputs_folder = Path(__file__).parent / "inputs" - - -class TestDebug(unittest.TestCase): - test_log1 = inputs_folder / "test.log" - - def test_logfile(self): - example_log = inputs_folder / "example.log" - test_log2 = inputs_folder / "test2.log" - - with MAD(stdout=self.test_log1, debug=True, raise_on_madng_error=False) as mad: - pass - time.sleep(0.1) # Wait for file to be written - with self.test_log1.open() as f, example_log.open() as f2: - self.assertEqual(f.read(), f2.read()) - - with MAD(stdout=test_log2, debug=True) as mad: - mad.send("!This is a line that does nothing") - mad.send("print('hello world')") - - with test_log2.open() as f: - text = f.read() - self.assertTrue("[!This is a line that does nothing]" in text) - self.assertTrue("hello world\n" in text) - self.assertTrue("[print('hello world')]" in text) - - with MAD(stdout=test_log2, debug=False) as mad: - mad.send("!This is a line that does nothing") - mad.send("print('hello world')") - - with test_log2.open() as f: - self.assertEqual(f.read(), "hello world\n") - - self.test_log1.unlink() - test_log2.unlink() - - def test_err(self): - # Run debug without stderr redirection - with ( - self.test_log1.open("w") as f, - MAD(debug=True, stdout=f, raise_on_madng_error=False) as mad, - ): - mad.psend("a = nil/2") - # receive the error before closing the pipe - self.assertRaises(RuntimeError, mad.recv) - with self.test_log1.open() as f: - # Check command was sent - file_text = f.read() - self.assertTrue("[py:__err(true); a = nil/2; py:__err(false);]" in file_text) - # Check error was not in stdout - self.assertFalse("***pymad.run:" in file_text) - - # Run debug with stderr redirection - with MAD(stdout=self.test_log1, redirect_stderr=True) as mad: - mad.psend("a = nil/2") - # receive the error before closing the pipe - self.assertRaises(RuntimeError, mad.recv) - with self.test_log1.open() as f: - # Check command was sent - file_text = f.read() - # Check error was in stdout - self.assertEqual("***pymad.run:", file_text[:13]) - self.test_log1.unlink() - - -if __name__ == "__main__": - unittest.main() + process = make_test_process() + process.send.side_effect = BrokenPipeError + process.stdout_file = MagicMock() + process.mad_read_stream.closed = False + process.mad_input_stream.closed = False + MadProcess.close(process) diff --git a/tests/test_examples_behaviour.py b/tests/test_examples_behaviour.py new file mode 100644 index 0000000..0929e08 --- /dev/null +++ b/tests/test_examples_behaviour.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import json +import os + +import numpy as np +import pytest + +from pymadng import MAD + + +def test_manual_close_without_context_manager(): + mad = MAD() + try: + mad.send("py:send(math.sqrt(16))") + assert mad.recv() == 4 + assert mad._MAD__process.process.poll() is None + finally: + mad.close() + + assert mad._MAD__process.process.poll() is not None + assert mad._MAD__process.mad_read_stream.closed + assert mad._MAD__process.mad_input_stream.closed + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="requires os.fork") +def test_forked_processes_have_independent_mad_sessions(): + read_fd, write_fd = os.pipe() + pid = os.fork() + + if pid == 0: + os.close(read_fd) + try: + with MAD() as child_mad: + child_mad["value"] = 7 + payload = {"value": int(child_mad["value"])} + except RuntimeError as exc: # pragma: no cover - child failure path + payload = {"error": repr(exc)} + finally: + os.write(write_fd, json.dumps(payload).encode("utf-8")) + os.close(write_fd) + os._exit(0) + + os.close(write_fd) + try: + with MAD() as parent_mad: + parent_mad["value"] = 3 + child_payload = os.read(read_fd, 4096).decode("utf-8") + pid_done, status = os.waitpid(pid, 0) + assert pid_done == pid + assert os.WIFEXITED(status) + assert os.WEXITSTATUS(status) == 0 + assert parent_mad["value"] == 3 + child_result = json.loads(child_payload) + assert "error" not in child_result + assert child_result["value"] == 7 + finally: + os.close(read_fd) + + +def test_multi_assignment_of_numpy_arrays(): + with MAD() as mad: + arr = np.arange(9, dtype=np.float64).reshape(3, 3) + mad["arr1", "arr2", "arr3"] = arr, arr * 2, arr * 3 + + recv1, recv2, recv3 = mad["arr1", "arr2", "arr3"] + + assert np.array_equal(recv1, arr) + assert np.array_equal(recv2, arr * 2) + assert np.array_equal(recv3, arr * 3) diff --git a/tests/test_io_and_loading.py b/tests/test_io_and_loading.py index 57e0d07..3dd75c1 100644 --- a/tests/test_io_and_loading.py +++ b/tests/test_io_and_loading.py @@ -1,65 +1,61 @@ -import unittest -from pathlib import Path +from __future__ import annotations import numpy as np from pymadng import MAD +A_MATRIX = np.arange(1, 21).reshape(4, 5) +B_MATRIX = np.arange(1, 7).reshape(2, 3) +C_MATRIX = (np.arange(1, 16) + 1j).reshape(5, 3) -class TestLoad(unittest.TestCase): - a = np.arange(1, 21).reshape(4, 5) - b = np.arange(1, 7).reshape(2, 3) - c = (np.arange(1, 16) + 1j).reshape(5, 3) - def test_load(self): - with MAD() as mad: - mad.load("MAD", "matrix") - self.assertTrue(mad.send("py:send(matrix == MAD.matrix)").recv()) +def test_load(tmp_path): + with MAD() as mad: + mad.load("MAD", "matrix") + assert mad.send("py:send(matrix == MAD.matrix)").recv() - mad.load("MAD.gmath") - self.assertTrue(mad.send("py:send(sin == MAD.gmath.sin)").recv()) - self.assertTrue(mad.send("py:send(cos == MAD.gmath.cos)").recv()) - self.assertTrue(mad.send("py:send(tan == MAD.gmath.tan)").recv()) + mad.load("MAD.gmath") + assert mad.send("py:send(sin == MAD.gmath.sin)").recv() + assert mad.send("py:send(cos == MAD.gmath.cos)").recv() + assert mad.send("py:send(tan == MAD.gmath.tan)").recv() - mad.load("MAD.element", "quadrupole", "sextupole", "drift") - self.assertTrue(mad.send("py:send(quadrupole == MAD.element.quadrupole)").recv()) - self.assertTrue(mad.send("py:send(sextupole == MAD.element.sextupole )").recv()) - self.assertTrue(mad.send("py:send(drift == MAD.element.drift )").recv()) + mad.load("MAD.element", "quadrupole", "sextupole", "drift") + assert mad.send("py:send(quadrupole == MAD.element.quadrupole)").recv() + assert mad.send("py:send(sextupole == MAD.element.sextupole )").recv() + assert mad.send("py:send(drift == MAD.element.drift )").recv() - test_file = Path("test.mad") - with test_file.open("w") as f: - f.write(""" + test_file = tmp_path / "test.mad" + test_file.write_text( + """ local matrix, cmatrix in MAD a = matrix(4, 5):seq() b = cmatrix(2, 3):seq() - """) - with MAD(stdout="/dev/null", redirect_stderr=True) as mad: - mad.loadfile("test.mad") - self.assertIsNone(mad.matrix) - self.assertTrue(np.all(mad.a == self.a)) - self.assertTrue(np.all(mad.b == self.b)) - test_file.unlink() - with test_file.open("w") as f: - f.write(""" + """ + ) + with MAD(stdout="/dev/null", redirect_stderr=True) as mad: + mad.loadfile(test_file) + assert mad.matrix is None + assert np.all(mad.a == A_MATRIX) + assert np.all(mad.b == B_MATRIX) + + test_file.write_text( + """ local matrix, cmatrix in MAD local a = matrix(4, 5):seq() local c = cmatrix(5, 3):seq() + 1i return {res1 = a * c, res2 = a * c:conj()} - """) - with MAD() as mad: - mad.loadfile("test.mad", "res1", "res2") - self.assertTrue(np.all(mad.res1 == np.matmul(self.a, self.c))) - self.assertTrue(np.all(mad.res2 == np.matmul(self.a, self.c.conj()))) - test_file.unlink() - - def test_globals(self): - with MAD() as mad: - mad.send_vars(a=1, b=2, c=3) - global_vars = mad.globals() - self.assertIn("a", global_vars) - self.assertIn("b", global_vars) - self.assertIn("c", global_vars) - - -if __name__ == "__main__": - unittest.main() + """ + ) + with MAD() as mad: + mad.loadfile(test_file, "res1", "res2") + assert np.all(mad.res1 == np.matmul(A_MATRIX, C_MATRIX)) + assert np.all(mad.res2 == np.matmul(A_MATRIX, C_MATRIX.conj())) + + +def test_globals(): + with MAD() as mad: + mad.send_vars(a=1, b=2, c=3) + global_vars = mad.globals() + assert "a" in global_vars + assert "b" in global_vars + assert "c" in global_vars diff --git a/tests/test_misc_types.py b/tests/test_misc_types.py index f109550..9224033 100644 --- a/tests/test_misc_types.py +++ b/tests/test_misc_types.py @@ -1,193 +1,226 @@ -import unittest -from pathlib import Path +from __future__ import annotations import numpy as np +import pytest from pymadng import MAD -inputs_folder = Path(__file__).parent / "inputs" -# TODO: Test the following functions: -# - eval -# - error on stdout = something strange +RANGE_CASES = [ + pytest.param( + "MAD.range(3, 11, 2)", + range(3, 12, 2), + False, + id="integer-range", + ), + pytest.param( + "MAD.nrange(3.5, 21.4, 12)", + np.linspace(3.5, 21.4, 12), + False, + id="numeric-range", + ), + pytest.param( + "MAD.nlogrange(1, 20, 20)", + np.geomspace(1, 20, 20), + False, + id="log-range", + ), + pytest.param( + "MAD.range(3, 11, 2)", + list(range(3, 12, 2)), + True, + id="integer-range-table", + ), + pytest.param( + "MAD.nrange(3.5, 21.4, 12)", + np.linspace(3.5, 21.4, 12), + True, + id="numeric-range-table", + ), + pytest.param( + "MAD.nlogrange(1, 20, 20)", + np.geomspace(1, 20, 20), + True, + id="log-range-table", + ), +] +SEND_RANGE_CASES = [ + pytest.param( + "send", + (range(3, 10, 1),), + "py:recv() + 1", + list(range(4, 12, 1)), + id="python-range", + ), + pytest.param( + "send_range", + (3.5, 21.4, 14), + "py:recv() + 2", + np.linspace(5.5, 23.4, 14), + id="numeric-range", + ), + pytest.param( + "send_logrange", + (1, 20, 20), + "py:recv()", + np.geomspace(1, 20, 20), + id="log-range", + ), +] -class TestNil(unittest.TestCase): - def test_send_recv(self): - with MAD() as mad: - mad.send(""" +TPSA_MONOMIALS = np.asarray( + [[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1], [2, 0, 0], [1, 1, 0]], + dtype=np.uint8, +) +TPSA_INDEX = ["000", "100", "010", "001", "200", "110"] + + +def assert_received_value(actual, expected): + if isinstance(expected, np.ndarray): + assert np.allclose(actual, expected) + else: + assert actual == expected + + +def test_send_recv_nil(): + with MAD() as mad: + mad.send( + """ local myNil = py:recv() py:send(myNil) py:send(nil) py:send() - """) - mad.send(None) - self.assertIsNone(mad.recv()) - self.assertIsNone(mad.recv()) - self.assertIsNone(mad.recv()) - - -class TestRngs(unittest.TestCase): - def test_recv(self): - with MAD() as mad: - mad.send(""" - irng = MAD.range(3, 11, 2) - rng = MAD.nrange(3.5, 21.4, 12) - lrng = MAD.nlogrange(1, 20, 20) - py:send(irng) - py:send(rng) - py:send(lrng) - py:send(irng:totable(), true) - py:send(rng:totable(), true) - py:send(lrng:totable(), true) - """) - - self.assertEqual( - mad.recv(), range(3, 12, 2) - ) # MAD is inclusive, python is exclusive (on stop) - self.assertTrue(np.allclose(mad.recv(), np.linspace(3.5, 21.4, 12))) - self.assertTrue(np.allclose(mad.recv(), np.geomspace(1, 20, 20))) - self.assertEqual( - mad.recv(), list(range(3, 12, 2)) - ) # MAD is inclusive, python is exclusive (on stop) - self.assertTrue(np.allclose(mad.recv(), np.linspace(3.5, 21.4, 12))) - self.assertTrue(np.allclose(mad.recv(), np.geomspace(1, 20, 20))) - - def test_send(self): - with MAD() as mad: - mad.send(""" - irng = py:recv() + 1 - rng = py:recv() + 2 - lrng = py:recv() - py:send(irng:totable(), true) - py:send(rng:totable(), true) - py:send(lrng:totable(), true) - """) - mad.send(range(3, 10, 1)) - mad.send_range(3.5, 21.4, 14) - mad.send_logrange(1, 20, 20) - self.assertEqual(mad.recv(), list(range(4, 12, 1))) - self.assertTrue(np.allclose(mad.recv(), np.linspace(5.5, 23.4, 14))) - self.assertTrue(np.allclose(mad.recv(), np.geomspace(1, 20, 20))) - - -class TestBool(unittest.TestCase): - def test_send_recv(self): - with MAD() as mad: - mad.send(""" - bool1 = py:recv() - bool2 = py:recv() - py:send(not bool1) - py:send(not bool2) - """) - mad.send(data=True) - mad.send(data=False) - self.assertEqual(mad.recv(), False) - self.assertEqual(mad.recv(), True) - - -class TestMono(unittest.TestCase): - def test_recv(self): - with MAD() as mad: - mad.send(""" - local m = MAD.monomial({1, 2, 3, 4, 6, 20, 100}) - py:send(m) - local m = MAD.monomial("WZ346oy") + """ + ) + mad.send(None) + assert mad.recv() is None + assert mad.recv() is None + assert mad.recv() is None + + +@pytest.mark.parametrize(("lua_expr", "expected", "as_table"), RANGE_CASES) +def test_recv_ranges(lua_expr, expected, as_table): + with MAD() as mad: + mad.send( + f""" + local value = {lua_expr} + py:send({"value:totable()" if as_table else "value"}{", true" if as_table else ""}) + """ + ) + assert_received_value(mad.recv(), expected) + + +@pytest.mark.parametrize(("send_method", "args", "lua_expr", "expected"), SEND_RANGE_CASES) +def test_send_ranges(send_method, args, lua_expr, expected): + with MAD() as mad: + mad.send( + f""" + local value = {lua_expr} + py:send(value:totable(), true) + """ + ) + getattr(mad, send_method)(*args) + assert_received_value(mad.recv(), expected) + + +@pytest.mark.parametrize(("value", "expected"), [(True, False), (False, True)]) +def test_send_recv_bool(value, expected): + with MAD() as mad: + mad.send( + """ + local value = py:recv() + py:send(not value) + """ + ) + mad.send(data=value) + assert mad.recv() == expected + + +@pytest.mark.parametrize( + ("lua_input", "expected"), + [ + ("{1, 2, 3, 4, 6, 20, 100}", [1, 2, 3, 4, 6, 20, 100]), + ('"WZ346oy"', [32, 35, 3, 4, 6, 50, 60]), + ], + ids=["from-table", "from-string"], +) +def test_recv_monomial(lua_input, expected): + with MAD() as mad: + mad.send( + f""" + local m = MAD.monomial({lua_input}) py:send(m) - """) + """ + ) + assert np.all(mad.recv() == expected) - self.assertTrue(np.all(mad.recv() == [1, 2, 3, 4, 6, 20, 100])) - self.assertTrue(np.all(mad.recv() == [32, 35, 3, 4, 6, 50, 60])) - def test_send_recv(self): - with MAD() as mad: - mad.send(""" +def test_send_recv_monomial(): + with MAD() as mad: + mad.send( + """ local m1 = py:recv() local m2 = py:recv() py:send(m1 + m2) - """) - rng = np.random.default_rng() - pym1 = rng.integers(0, 255, 20, dtype=np.ubyte) - pym2 = rng.integers(0, 255, 20, dtype=np.ubyte) - mad.send(pym1) - mad.send(pym2) - mad_res = mad.recv() - self.assertTrue(np.all(mad_res == pym1 + pym2)) - # Check the return is a monomial - self.assertEqual(mad_res.dtype, np.dtype("ubyte")) - - -class TestTPSA(unittest.TestCase): - def test_recv_real(self): - with MAD() as mad: - mad.send(""" - local d = MAD.gtpsad(3, 6) - res = MAD.tpsa(6):set(1,2):set(2, 1) - res2 = res:copy():set(3, 1) - res3 = res2:copy():set(4, 1) - py:send(res:axypbzpc(res2, res3, 1, 2, 3)) - """) - monomials, coefficients = mad.recv() - self.assertTrue(np.all(monomials[0] == [0, 0, 0])) - self.assertTrue(np.all(monomials[1] == [1, 0, 0])) - self.assertTrue(np.all(monomials[2] == [0, 1, 0])) - self.assertTrue(np.all(monomials[3] == [0, 0, 1])) - self.assertTrue(np.all(monomials[4] == [2, 0, 0])) - self.assertTrue(np.all(monomials[5] == [1, 1, 0])) - self.assertTrue(np.all(coefficients == [11, 6, 4, 2, 1, 1])) - - def test_recv_cpx(self): - with MAD() as mad: - mad.send(""" + """ + ) + rng = np.random.default_rng() + pym1 = rng.integers(0, 255, 20, dtype=np.ubyte) + pym2 = rng.integers(0, 255, 20, dtype=np.ubyte) + mad.send(pym1) + mad.send(pym2) + mad_res = mad.recv() + assert np.all(mad_res == pym1 + pym2) + assert mad_res.dtype == np.dtype("ubyte") + + +@pytest.mark.parametrize( + ("lua_type", "setter_values", "expected_coefficients"), + [ + pytest.param( + "tpsa", + ("2", "1", "1", "1"), + [11, 6, 4, 2, 1, 1], + id="real", + ), + pytest.param( + "ctpsa", + ("2+1i", "1+2i", "1+2i", "1+2i"), + [10 + 6j, 2 + 14j, 2 + 9j, 2 + 4j, -3 + 4j, -3 + 4j], + id="complex", + ), + ], +) +def test_recv_tpsa(lua_type, setter_values, expected_coefficients): + with MAD() as mad: + first, second, third, fourth = setter_values + mad.send( + f""" local d = MAD.gtpsad(3, 6) - res = MAD.ctpsa(6):set(1,2+1i):set(2, 1+2i) - res2 = res:copy():set(3, 1+2i) - res3 = res2:copy():set(4, 1+2i) + res = MAD.{lua_type}(6):set(1,{first}):set(2, {second}) + res2 = res:copy():set(3, {third}) + res3 = res2:copy():set(4, {fourth}) py:send(res:axypbzpc(res2, res3, 1, 2, 3)) - """) - monomials, coefficients = mad.recv() - self.assertTrue(np.all(monomials[0] == [0, 0, 0])) - self.assertTrue(np.all(monomials[1] == [1, 0, 0])) - self.assertTrue(np.all(monomials[2] == [0, 1, 0])) - self.assertTrue(np.all(monomials[3] == [0, 0, 1])) - self.assertTrue(np.all(monomials[4] == [2, 0, 0])) - self.assertTrue(np.all(monomials[5] == [1, 1, 0])) - self.assertTrue( - np.all(coefficients == [10 + 6j, 2 + 14j, 2 + 9j, 2 + 4j, -3 + 4j, -3 + 4j]) - ) - - # Might be worth checking if the tab can be converted into a tpsa from Monomials. (jgray 2025) - def test_send_tpsa(self): - with MAD() as mad: - mad.send(""" - tab = py:recv() - index_part = {} - for i = 1, #tab do - index_part[i] = tab[i] - end - py:send(tab, true) - py:send(index_part, true) - """) - monos = np.asarray( - [[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1], [2, 0, 0], [1, 1, 0]], dtype=np.uint8 - ) - coefficients = [11, 6, 4, 2, 1, 1] - expected_index = ["000", "100", "010", "001", "200", "110"] - - mad.send_tpsa(monos, coefficients) - - whole_tab = mad.recv("tab") - index_part = mad.recv("index_part") - self.assertEqual(index_part, expected_index) - for key, value in whole_tab.items(): - if isinstance(key, int): - self.assertTrue(value == expected_index[key - 1]) - else: - idx = expected_index.index(key) - self.assertTrue(whole_tab[key] == coefficients[idx]) - - def test_send_ctpsa(self): - with MAD() as mad: - mad.send(""" + """ + ) + monomials, coefficients = mad.recv() + assert np.array_equal(monomials, TPSA_MONOMIALS) + assert np.all(coefficients == expected_coefficients) + + +@pytest.mark.parametrize( + ("send_method", "coefficients"), + [ + pytest.param("send_tpsa", [11, 6, 4, 2, 1, 1], id="real"), + pytest.param( + "send_cpx_tpsa", [10 + 6j, 2 + 14j, 2 + 9j, 2 + 4j, -3 + 4j, -3 + 4j], id="complex" + ), + ], +) +def test_send_tpsa(send_method, coefficients): + with MAD() as mad: + mad.send( + """ tab = py:recv() index_part = {} for i = 1, #tab do @@ -195,28 +228,25 @@ def test_send_ctpsa(self): end py:send(tab, true) py:send(index_part, true) - """) - monos = np.asarray( - [[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1], [2, 0, 0], [1, 1, 0]], dtype=np.uint8 - ) - coefficients = [10 + 6j, 2 + 14j, 2 + 9j, 2 + 4j, -3 + 4j, -3 + 4j] - expected_index = ["000", "100", "010", "001", "200", "110"] - - mad.send_cpx_tpsa(monos, coefficients) - - whole_tab = mad.recv("tab") - index_part = mad.recv("index_part") - self.assertEqual(index_part, expected_index) - for key, value in whole_tab.items(): - if isinstance(key, int): - self.assertTrue(value == expected_index[key - 1]) - else: - idx = expected_index.index(key) - self.assertTrue(whole_tab[key] == coefficients[idx]) - - def test_send_recv_damap(self): - with MAD() as mad: - mad.send(""" + """ + ) + getattr(mad, send_method)(TPSA_MONOMIALS, coefficients) + + whole_tab = mad.recv("tab") + index_part = mad.recv("index_part") + assert index_part == TPSA_INDEX + for key, value in whole_tab.items(): + if isinstance(key, int): + assert value == TPSA_INDEX[key - 1] + else: + idx = TPSA_INDEX.index(key) + assert whole_tab[key] == coefficients[idx] + + +def test_send_recv_damap(): + with MAD() as mad: + mad.send( + """ py:__err(true) local sin in MAD.gmath MAD.gtpsad(6, 5) @@ -227,13 +257,10 @@ def test_send_recv_damap(self): py:send(res) recved = MAD.tpsa():fromtable(py:recv()) py:send(recved) - """) - init = mad.recv() - mad.send_tpsa(*init) - final = mad.recv() - self.assertTrue((init[0] == final[0]).all()) - self.assertTrue((init[1] == final[1]).all()) - - -if __name__ == "__main__": - unittest.main() + """ + ) + init = mad.recv() + mad.send_tpsa(*init) + final = mad.recv() + assert (init[0] == final[0]).all() + assert (init[1] == final[1]).all() diff --git a/tests/test_numeric_types.py b/tests/test_numeric_types.py index 82a0126..2b9e63f 100644 --- a/tests/test_numeric_types.py +++ b/tests/test_numeric_types.py @@ -1,156 +1,197 @@ -import unittest +from __future__ import annotations import numpy as np +import pytest from pymadng import MAD - -class TestList(unittest.TestCase): - def test_send_recv(self): - with MAD() as mad: - test_list = [[1, 2, 3, 4, 5, 6, 7, 8, 9]] * 2 - mad.send(""" +FLT_EPS = 2**-52 +FLT_TINY = 2**-1022 +FLT_HUGE = 2**1023 +FLOAT_VALUES = [ + 0, + FLT_TINY, + 2**-64, + 2**-63, + 2**-53, + FLT_EPS, + 2**-52, + 2 * FLT_EPS, + 2**-32, + 2**-31, + 1e-9, + 0.1 - FLT_EPS, + 0.1, + 0.1 + FLT_EPS, + 0.5, + 0.7 - FLT_EPS, + 0.7, + 0.7 + FLT_EPS, + 1 - FLT_EPS, + 1, + 1 + FLT_EPS, + 1.1, + 1.7, + 2, + 10, + 1e2, + 1e3, + 1e6, + 1e9, + 2**31, + 2**32, + 2**52, + 2**53, + 2**63, + 2**64, + FLT_HUGE, +] +INT_VALUES = [0, 1, 2, 10, 1e2, 1e3, 1e6, 1e9, 2**31 - 1] + + +def test_send_recv_list(): + with MAD() as mad: + test_list = [[1, 2, 3, 4, 5, 6, 7, 8, 9]] * 2 + mad.send( + """ list = py:recv() list[1][1] = 10 list[2][1] = 10 py:send(list) - """) - mad.send(test_list) - test_list[0][0] = 10 - test_list[1][0] = 10 - mad_list = mad.recv("list") - for i, inner_list in enumerate(mad_list): - for j, val in enumerate(inner_list): - self.assertEqual( - val, - test_list[i][j], - f"Mismatch at index [{i}][{j}]: {val} != {test_list[i][j]}", - ) - - def test_send_recv_wref(self): - with MAD() as mad: - python_dict = {1: 1, 2: 2, 3: 3, 4: 4, 5: 5, "a": 10, "b": 3, "c": 4} - mad.send(""" + """ + ) + mad.send(test_list) + test_list[0][0] = 10 + test_list[1][0] = 10 + mad_list = mad.recv("list") + + for i, inner_list in enumerate(mad_list): + for j, value in enumerate(inner_list): + assert value == test_list[i][j], ( + f"Mismatch at index [{i}][{j}]: {value} != {test_list[i][j]}" + ) + + +def test_send_recv_list_with_ref(): + with MAD() as mad: + python_dict = {1: 1, 2: 2, 3: 3, 4: 4, 5: 5, "a": 10, "b": 3, "c": 4} + mad.send( + """ list = {MAD.object "a" {a = 2}, MAD.object "b" {b = 6}} list2 = py:recv() py:send(list) py:send(list2) - """).send(python_dict) - list1 = mad.recv("list") - list2 = mad.recv("list2") - self.assertEqual(len(list1), 2) - - self.assertEqual(list2.eval().keys(), python_dict.keys()) - self.assertEqual(sorted(list2.eval().values()), sorted(python_dict.values())) - - self.assertEqual(list1[0].a, 2) - self.assertEqual(list1[1].b, 6) - - -class TestNums(unittest.TestCase): - eps = 2**-52 - tiny = 2**-1022 - huge = 2**1023 - # fmt: off - flt_lst = [ - 0, tiny, 2**-64, 2**-63, 2**-53, eps, 2**-52, 2*eps, 2**-32, 2**-31, 1e-9, - 0.1-eps, 0.1, 0.1+eps, 0.5, 0.7-eps, 0.7, 0.7+eps, 1-eps, 1, 1+eps, - 1.1, 1.7, 2, 10, 1e2, 1e3, 1e6, 1e9, 2**31, 2**32, 2**52, 2**53, - 2**63, 2**64, huge - ] - # fmt: on - - def test_send_recv_int(self): - with MAD() as mad: - int_lst = [0, 1, 2, 10, 1e2, 1e3, 1e6, 1e9, 2**31 - 1] - for i in range(len(int_lst)): - mad.send(""" - local is_integer in MAD.typeid - local num = py:recv() - py:send( num) - py:send(-num) - py:send(is_integer(num)) - """) - mad.send(int_lst[i]) - recv_num = mad.recv() - self.assertEqual(recv_num, int_lst[i]) - self.assertTrue(isinstance(recv_num, np.int32)) - recv_num = mad.recv() - self.assertEqual(recv_num, -int_lst[i]) - self.assertTrue(isinstance(recv_num, np.int32)) - self.assertTrue(mad.recv()) - - def test_send_recv_num(self): - with MAD() as mad: - for i in range(len(self.flt_lst)): - mad.send(""" - local num = py:recv() - local negative = py:recv() - py:send(num) - py:send(negative) - py:send(num * 1.61) - """) - mad.send(self.flt_lst[i]) - mad.send(-self.flt_lst[i]) - self.assertEqual(mad.recv(), self.flt_lst[i]) # Check individual floats - self.assertEqual(mad.recv(), -self.flt_lst[i]) # Check negation - self.assertEqual(mad.recv(), self.flt_lst[i] * 1.61) # Check manipulation - - def test_send_recv_cpx(self): - with MAD() as mad: - for i in range(len(self.flt_lst)): - for j in range(len(self.flt_lst)): - mad.send(""" + """ + ).send(python_dict) + list1 = mad.recv("list") + list2 = mad.recv("list2") + received_dict = list2.eval() + + assert len(list1) == 2 + assert received_dict == python_dict + assert list1[0].a == 2 + assert list1[1].b == 6 + + +@pytest.mark.parametrize("value", INT_VALUES) +def test_send_recv_int(value): + with MAD() as mad: + mad.send( + """ + local is_integer in MAD.typeid + local num = py:recv() + py:send(num) + py:send(-num) + py:send(is_integer(num)) + """ + ) + mad.send(value) + recv_num = mad.recv() + assert recv_num == value + assert isinstance(recv_num, np.int32) + + recv_num = mad.recv() + assert recv_num == -value + assert isinstance(recv_num, np.int32) + assert mad.recv() + + +@pytest.mark.parametrize("value", FLOAT_VALUES) +def test_send_recv_num(value): + with MAD() as mad: + mad.send( + """ + local num = py:recv() + local negative = py:recv() + py:send(num) + py:send(negative) + py:send(num * 1.61) + """ + ) + mad.send(value) + mad.send(-value) + assert mad.recv() == value + assert mad.recv() == -value + assert mad.recv() == value * 1.61 + + +def test_send_recv_cpx(): + with MAD() as mad: + for real in FLOAT_VALUES: + for imag in FLOAT_VALUES: + mad.send( + """ local my_cpx = py:recv() py:send(my_cpx) py:send(-my_cpx) py:send(my_cpx * 1.31i) - """) - my_cpx = self.flt_lst[i] + 1j * self.flt_lst[j] - mad.send(my_cpx) - self.assertEqual(mad.recv(), my_cpx) - self.assertEqual(mad.recv(), -my_cpx) - self.assertEqual(mad.recv(), my_cpx * 1.31j) - - -class TestMatrices(unittest.TestCase): - def test_send_recv_imat(self): - with MAD() as mad: - mad.send(""" - local imat = py:recv() - py:send(imat) - py:send(MAD.imatrix(3, 5):seq()) - """) - rng = np.random.default_rng() - imat = rng.integers(0, 255, (5, 5), dtype=np.int32) - mad.send(imat) - self.assertTrue(np.all(mad.recv() == imat)) - self.assertTrue(np.all(mad.recv() == np.arange(1, 16).reshape(3, 5))) - - def test_send_recv_mat(self): - with MAD() as mad: - mad.send(""" - local mat = py:recv() - py:send(mat) - py:send(MAD.matrix(3, 5):seq() / 2) - """) - mat = np.arange(1, 25).reshape(4, 6) / 4 - mad.send(mat) - self.assertTrue(np.all(mad.recv() == mat)) - self.assertTrue(np.all(mad.recv() == np.arange(1, 16).reshape(3, 5) / 2)) - - def test_send_recv_cmat(self): - with MAD() as mad: - mad.send(""" - local cmat = py:recv() - py:send(cmat) - py:send(MAD.cmatrix(3, 5):seq() / 2i) - """) - cmat = np.arange(1, 25).reshape(4, 6) / 4 + 1j * np.arange(1, 25).reshape(4, 6) / 4 - mad.send(cmat) - self.assertTrue(np.all(mad.recv() == cmat)) - self.assertTrue(np.all(mad.recv() == (np.arange(1, 16).reshape(3, 5) / 2j))) - - -if __name__ == "__main__": - unittest.main() + """ + ) + value = real + 1j * imag + mad.send(value) + assert mad.recv() == value + assert mad.recv() == -value + assert mad.recv() == value * 1.31j + + +@pytest.mark.parametrize( + ("lua_ctor", "array"), + [ + ( + "MAD.imatrix(3, 5):seq()", + np.random.default_rng().integers(0, 255, (5, 5), dtype=np.int32), + ), + ("MAD.matrix(3, 5):seq() / 2", np.arange(1, 25).reshape(4, 6) / 4), + ( + "MAD.cmatrix(3, 5):seq() / 2i", + np.arange(1, 25).reshape(4, 6) / 4 + 1j * np.arange(1, 25).reshape(4, 6) / 4, + ), + ], + ids=["imat", "mat", "cmat"], +) +def test_send_recv_matrices(lua_ctor, array): + with MAD() as mad: + local_name = ( + "imat" + if np.issubdtype(array.dtype, np.integer) + else "cmat" + if np.iscomplexobj(array) + else "mat" + ) + mad.send( + f""" + local {local_name} = py:recv() + py:send({local_name}) + py:send({lua_ctor}) + """ + ) + mad.send(array) + assert np.all(mad.recv() == array) + + expected = np.arange(1, 16).reshape(3, 5) + if np.issubdtype(array.dtype, np.integer): + assert np.all(mad.recv() == expected) + elif np.iscomplexobj(array): + assert np.all(mad.recv() == (expected / 2j)) + else: + assert np.all(mad.recv() == (expected / 2)) diff --git a/tests/test_object_wrapping.py b/tests/test_object_wrapping.py index 8253a89..19813df 100644 --- a/tests/test_object_wrapping.py +++ b/tests/test_object_wrapping.py @@ -1,177 +1,207 @@ +from __future__ import annotations + +import copy import math import sys -import unittest -from pathlib import Path import numpy as np import pandas as pd +import pytest import tfs from pymadng import MAD from pymadng.madp_classes import MadLastRef, MadRef -# TODO: Test the following functions: -# - __str__ on mad references (low priority) - - -class TestGetSet(unittest.TestCase): - def test_get(self): - with MAD(stdout="/dev/null", redirect_stderr=True) as mad: - mad.load("element", "quadrupole") - self.assertEqual(mad.asdfg, None) - mad.send("""qd = quadrupole {knl={0, 0.25}, l = 1}""") - mad.send("""qf = quadrupole {qd = qd}""") - qd, qf = mad["qd", "qf"] - self.assertEqual(qd._name, "qd") - self.assertEqual(qd._parent, None) - self.assertEqual(qd._mad, mad._MAD__process) - self.assertEqual(qd.knl.eval(), [0, 0.25]) - self.assertEqual(qd.l, 1) - self.assertRaises(AttributeError, lambda: qd.asdfg) - self.assertRaises(KeyError, lambda: qd["asdfg"]) - self.assertRaises(IndexError, lambda: qd[1]) - self.assertTrue(isinstance(qf.qd, MadRef)) - self.assertEqual(qf.qd.knl.eval(), [0, 0.25]) - self.assertEqual(qf.qd.l, 1) - self.assertEqual(qf.qd, qd) - - mad.send("objList = {qd, qf, qd, qf, qd} py:send(objList)") - obj_lst = mad.recv("objList") - for i in range(len(obj_lst)): - if i % 2 != 0: - self.assertTrue(isinstance(obj_lst[i].qd, MadRef)) - self.assertEqual(obj_lst[i].qd._parent, f"objList[{i + 1}]") - self.assertEqual(obj_lst[i].qd.knl.eval(), [0, 0.25]) - self.assertEqual(obj_lst[i].qd.l, 1) - self.assertEqual(obj_lst[i].qd, qd) - - else: - self.assertEqual(obj_lst[i].knl.eval(), [0, 0.25]) - self.assertEqual(obj_lst[i].l, 1) - self.assertEqual(obj_lst[i], qd) - self.assertEqual(obj_lst[i]._parent, "objList") - - def test_set(self): # Need more? - with MAD() as mad: - mad.load("element", "quadrupole") - mad.send("""qd = quadrupole {knl={0, 0.25}, l = 1} py:send(qd)""") - mad["qd2"] = mad.recv("qd") - self.assertEqual(mad.qd2._name, "qd2") - self.assertEqual(mad.qd2._parent, None) - self.assertEqual(mad.qd2._mad, mad._MAD__process) - self.assertEqual(mad.qd2.knl.eval(), [0, 0.25]) - self.assertEqual(mad.qd2.l, 1) - self.assertEqual(mad.qd2, mad.qd) - mad["a", "b"] = mad.MAD.gmath.reim(9.75 + 1.5j) - self.assertEqual(mad.a, 9.75) - self.assertEqual(mad.b, 1.5) - mad.send("f = \\-> (1, 2, 3, 9, 8, 7)") - mad["r1", "r2", "r3", "r4", "r5", "r6"] = mad.f() - self.assertEqual(mad.r1, 1) - self.assertEqual(mad.r2, 2) - self.assertEqual(mad.r3, 3) - self.assertEqual(mad.r4, 9) - self.assertEqual(mad.r5, 8) - self.assertEqual(mad.r6, 7) - - def test_send_vars(self): - with MAD() as mad: - mad.send_vars(a=1, b=2.5, c="test", d=[1, 2, 3]) - self.assertEqual(mad.a, 1) - self.assertEqual(mad.b, 2.5) - self.assertEqual(mad.c, "test") - self.assertEqual(mad.d.eval(), [1, 2, 3]) - - def test_recv_vars(self): - with MAD() as mad: - mad.send_vars(a=1, b=2.5, c="test", d=[1, 2, 3]) - a, b, c, d = mad.recv_vars("a", "b", "c", "d") - self.assertEqual(a, 1) - self.assertEqual(b, 2.5) - self.assertEqual(c, "test") - self.assertEqual(d.eval(), [1, 2, 3]) - - def test_quote_strings(self): - with MAD() as mad: - self.assertEqual(mad.quote_strings("test"), "'test'") - self.assertEqual(mad.quote_strings(["a", "b", "c"]), ["'a'", "'b'", "'c'"]) - - def test_create_deferred_expression(self): - with MAD() as mad: - deferred = mad.create_deferred_expression(a="x + y", b="x * y") - mad.send_vars(x=2, y=3) - self.assertEqual(deferred.a, 5) # x + y = 2 + 3 - self.assertEqual(deferred.b, 6) # x * y = 2 * 3 - mad.send_vars(x=5, y=10) - self.assertEqual(deferred.a, 15) # x + y = 5 + 10 - self.assertEqual(deferred.b, 50) - - mad["deferred"] = mad.create_deferred_expression(a="x - y", b="x / y") - mad.send_vars(x=10, y=2) - self.assertEqual(mad.deferred.a, 8) # x - y = 10 - 2 - self.assertEqual(mad.deferred.b, 5) # x / y = 10 / 2 - - -class TestObjFun(unittest.TestCase): - def test_call_obj(self): - with MAD(py_name="python") as mad: - mad.load("element", "quadrupole", "sextupole") - qd = mad.quadrupole(knl=[0, 0.25], l=1) - sd = mad.sextupole(knl=[0, 0.25, 0.5], l=1) - - mad["qd"] = qd - self.assertEqual(mad.qd._name, "qd") - self.assertEqual(mad.qd._parent, None) - self.assertEqual(mad.qd._mad, mad._MAD__process) - self.assertEqual(mad.qd.knl.eval(), [0, 0.25]) - self.assertEqual(mad.qd.l, 1) - - sdc = sd - mad["sd"] = sd - del sd - self.assertEqual(mad.sd._name, "sd") - self.assertEqual(mad.sd._parent, None) - self.assertEqual(mad.sd._mad, mad._MAD__process) - self.assertEqual(mad.sd.knl.eval(), [0, 0.25, 0.5]) - self.assertEqual(mad.sd.l, 1) - - # Reference counting - qd = mad.quadrupole(knl=[0, 0.3], l=1) - self.assertEqual(sdc._name, "_last[2]") - self.assertEqual(qd._name, "_last[3]") - self.assertEqual(qd.knl.eval(), [0, 0.3]) - qd = mad.quadrupole(knl=[0, 0.25], l=1) - self.assertEqual(qd._name, "_last[1]") - - def test_call_last(self): - with MAD() as mad: - mad.send("func_test = \\a-> \\b-> \\c-> a+b*c") - self.assertRaises(TypeError, lambda: mad.MAD()) - self.assertEqual(mad.func_test(1)(2)(3), 7) - - def test_call_fail(self): - with MAD(stdout="/dev/null", redirect_stderr=True) as mad: - mad.send("func_test = \\a-> \\b-> \\c-> 'a'+b") - mad.func_test(1)(2)(3) - self.assertRaises(RuntimeError, lambda: mad.recv()) - self.assertRaises(RuntimeError, lambda: mad.mtable.read("'abad.tfs'").eval()) - - def test_call_func(self): - with MAD(py_name="python") as mad: - mad.load("element", "quadrupole") - mad["qd"] = mad.quadrupole(knl=[0, 0.25], l=1) - mad.qd.select() - mad["qdSelected"] = mad.qd.is_selected() - self.assertTrue(mad.qdSelected) - mad.qd.deselect() - mad["qdSelected"] = mad.qd.is_selected() - self.assertFalse(mad.qdSelected) - mad.qd.set_variables({"l": 2}) - self.assertEqual(mad.qd.l, 2) - - def test_mult_rtrn(self): - with MAD() as mad: - mad.send(""" + +def test_get(): + with MAD(stdout="/dev/null", redirect_stderr=True) as mad: + mad.load("element", "quadrupole") + assert mad.asdfg is None + mad.send("""qd = quadrupole {knl={0, 0.25}, l = 1}""") + mad.send("""qf = quadrupole {qd = qd}""") + qd, qf = mad["qd", "qf"] + assert qd._name == "qd" + assert qd._parent is None + assert qd._mad == mad._MAD__process + assert qd.knl.eval() == [0, 0.25] + assert qd.l == 1 + with pytest.raises(AttributeError): + qd.asdfg + with pytest.raises(KeyError): + qd["asdfg"] + with pytest.raises(IndexError): + qd[1] + assert isinstance(qf.qd, MadRef) + assert qf.qd.knl.eval() == [0, 0.25] + assert qf.qd.l == 1 + assert qf.qd == qd + + mad.send("objList = {qd, qf, qd, qf, qd} py:send(objList)") + obj_list = mad.recv("objList") + for i, item in enumerate(obj_list): + if i % 2 != 0: + assert isinstance(item.qd, MadRef) + assert item.qd._parent == f"objList[{i + 1}]" + assert item.qd.knl.eval() == [0, 0.25] + assert item.qd.l == 1 + assert item.qd == qd + else: + assert item.knl.eval() == [0, 0.25] + assert item.l == 1 + assert item == qd + assert item._parent == "objList" + + assert "knl" in dir(qd) + assert all(not name.startswith("_") for name in dir(qd)) + + +def test_deepcopy_warning(): + with MAD(stdout="/dev/null", redirect_stderr=True) as mad: + mad.load("element", "quadrupole") + mad.send("""qd = quadrupole {knl={0, 0.25}, l = 1}""") + qd = mad.qd + + with pytest.warns(UserWarning, match="An attempt to deepcopy a MadRef"): + qd_copy = copy.deepcopy(qd) + + assert qd_copy == qd + + +def test_invalid_reference_setitem_type(): + with MAD(stdout="/dev/null", redirect_stderr=True) as mad: + mad.load("element", "quadrupole") + mad.send("""qd = quadrupole {knl={0, 0.25}, l = 1}""") + + with pytest.raises(TypeError, match="expected string or int"): + mad.qd[1.5] = 2 + + +def test_set(): + with MAD() as mad: + mad.load("element", "quadrupole") + mad.send("""qd = quadrupole {knl={0, 0.25}, l = 1} py:send(qd)""") + mad["qd2"] = mad.recv("qd") + assert mad.qd2._name == "qd2" + assert mad.qd2._parent is None + assert mad.qd2._mad == mad._MAD__process + assert mad.qd2.knl.eval() == [0, 0.25] + assert mad.qd2.l == 1 + assert mad.qd2 == mad.qd + mad["a", "b"] = mad.MAD.gmath.reim(9.75 + 1.5j) + assert mad.a == 9.75 + assert mad.b == 1.5 + mad.send("f = \\-> (1, 2, 3, 9, 8, 7)") + mad["r1", "r2", "r3", "r4", "r5", "r6"] = mad.f() + assert (mad.r1, mad.r2, mad.r3, mad.r4, mad.r5, mad.r6) == (1, 2, 3, 9, 8, 7) + + +def test_send_vars(): + with MAD() as mad: + mad.send_vars(a=1, b=2.5, c="test", d=[1, 2, 3]) + assert mad.a == 1 + assert mad.b == 2.5 + assert mad.c == "test" + assert mad.d.eval() == [1, 2, 3] + + +def test_recv_vars(): + with MAD() as mad: + mad.send_vars(a=1, b=2.5, c="test", d=[1, 2, 3]) + a, b, c, d = mad.recv_vars("a", "b", "c", "d") + assert a == 1 + assert b == 2.5 + assert c == "test" + assert d.eval() == [1, 2, 3] + + +def test_quote_strings(): + with MAD() as mad: + assert mad.quote_strings("test") == "'test'" + assert mad.quote_strings(["a", "b", "c"]) == ["'a'", "'b'", "'c'"] + + +def test_create_deferred_expression(): + with MAD() as mad: + deferred = mad.create_deferred_expression(a="x + y", b="x * y") + mad.send_vars(x=2, y=3) + assert deferred.a == 5 + assert deferred.b == 6 + mad.send_vars(x=5, y=10) + assert deferred.a == 15 + assert deferred.b == 50 + + mad["deferred"] = mad.create_deferred_expression(a="x - y", b="x / y") + mad.send_vars(x=10, y=2) + assert mad.deferred.a == 8 + assert mad.deferred.b == 5 + + +def test_call_obj(): + with MAD(py_name="python") as mad: + mad.load("element", "quadrupole", "sextupole") + qd = mad.quadrupole(knl=[0, 0.25], l=1) + sd = mad.sextupole(knl=[0, 0.25, 0.5], l=1) + + mad["qd"] = qd + assert mad.qd._name == "qd" + assert mad.qd._parent is None + assert mad.qd._mad == mad._MAD__process + assert mad.qd.knl.eval() == [0, 0.25] + assert mad.qd.l == 1 + + sdc = sd + mad["sd"] = sd + del sd + assert mad.sd._name == "sd" + assert mad.sd._parent is None + assert mad.sd._mad == mad._MAD__process + assert mad.sd.knl.eval() == [0, 0.25, 0.5] + assert mad.sd.l == 1 + + qd = mad.quadrupole(knl=[0, 0.3], l=1) + assert sdc._name == "_last[2]" + assert qd._name == "_last[3]" + assert qd.knl.eval() == [0, 0.3] + qd = mad.quadrupole(knl=[0, 0.25], l=1) + assert qd._name == "_last[1]" + + +def test_call_last(): + with MAD() as mad: + mad.send("func_test = \\a-> \\b-> \\c-> a+b*c") + with pytest.raises(TypeError): + mad.MAD() + assert mad.func_test(1)(2)(3) == 7 + + +def test_call_fail(): + with MAD(stdout="/dev/null", redirect_stderr=True) as mad: + mad.send("func_test = \\a-> \\b-> \\c-> 'a'+b") + mad.func_test(1)(2)(3) + with pytest.raises(RuntimeError): + mad.recv() + with pytest.raises(RuntimeError): + mad.mtable.read("'abad.tfs'").eval() + + +def test_call_func(): + with MAD(py_name="python") as mad: + mad.load("element", "quadrupole") + mad["qd"] = mad.quadrupole(knl=[0, 0.25], l=1) + mad.qd.select() + mad["qdSelected"] = mad.qd.is_selected() + assert mad.qdSelected + mad.qd.deselect() + mad["qdSelected"] = mad.qd.is_selected() + assert not mad.qdSelected + mad.qd.set_variables({"l": 2}) + assert mad.qd.l == 2 + + +def test_mult_rtrn(): + with MAD() as mad: + mad.send( + """ obj = MAD.object "obj" {a = 1, b = 2} local function mult_rtrn () obj2 = MAD.object "obj2" {a = 2, b = 3} @@ -180,189 +210,181 @@ def test_mult_rtrn(self): last_rtn = __mklast__(mult_rtrn()) lastobj = __mklast__(obj) notLast = {mult_rtrn()} - """) - - mad["o11", "o12", "o13", "o2"] = mad._MAD__get_MadRef("last_rtn") - mad["p11", "p12", "p13", "p2"] = mad._MAD__get_MadRef("notLast") - mad["objCpy"] = mad._MAD__get_MadRef("lastobj") # Test single object in __mklast__ - self.assertEqual(mad.o11.a, 1) - self.assertEqual(mad.o11.b, 2) - self.assertEqual(mad.o12.a, 1) - self.assertEqual(mad.o12.b, 2) - self.assertEqual(mad.o13.a, 1) - self.assertEqual(mad.o13.b, 2) - self.assertEqual(mad.o2.a, 2) - self.assertEqual(mad.o2.b, 3) - self.assertEqual(mad.p11.a, 1) - self.assertEqual(mad.p11.b, 2) - self.assertEqual(mad.p12.a, 1) - self.assertEqual(mad.p12.b, 2) - self.assertEqual(mad.p13.a, 1) - self.assertEqual(mad.p13.b, 2) - self.assertEqual(mad.p2.a, 2) - self.assertEqual(mad.p2.b, 3) - self.assertEqual(mad.objCpy, mad.obj) - - def test_madx(self): - from pymadng import MAD - - mad = MAD() - self.assertEqual(mad.MADX.abs(-1), 1) - self.assertEqual(mad.MADX.ceil(1.2), 2) - - self.assertEqual(mad.MAD.MADX.abs(-1), 1) - self.assertEqual(mad.MAD.MADX.ceil(1.2), 2) - - self.assertEqual(mad.MADX.MAD.MADX.abs(-1), 1) - self.assertEqual(mad.MADX.MAD.MADX.ceil(1.2), 2) - seq_file = Path("test.seq") - with seq_file.open("w") as f: - f.write(""" + """ + ) + + mad["o11", "o12", "o13", "o2"] = mad._MAD__get_MadRef("last_rtn") + mad["p11", "p12", "p13", "p2"] = mad._MAD__get_MadRef("notLast") + mad["objCpy"] = mad._MAD__get_MadRef("lastobj") + assert mad.o11.a == 1 + assert mad.o11.b == 2 + assert mad.o12.a == 1 + assert mad.o12.b == 2 + assert mad.o13.a == 1 + assert mad.o13.b == 2 + assert mad.o2.a == 2 + assert mad.o2.b == 3 + assert mad.p11.a == 1 + assert mad.p11.b == 2 + assert mad.p12.a == 1 + assert mad.p12.b == 2 + assert mad.p13.a == 1 + assert mad.p13.b == 2 + assert mad.p2.a == 2 + assert mad.p2.b == 3 + assert mad.objCpy == mad.obj + + +def test_madx(tmp_path): + with MAD() as mad: + assert mad.MADX.abs(-1) == 1 + assert mad.MADX.ceil(1.2) == 2 + + assert mad.MAD.MADX.abs(-1) == 1 + assert mad.MAD.MADX.ceil(1.2) == 2 + + assert mad.MADX.MAD.MADX.abs(-1) == 1 + assert mad.MADX.MAD.MADX.ceil(1.2) == 2 + seq_file = tmp_path / "test.seq" + seq_file.write_text( + """ qd: quadrupole, l=1, knl:={0, 0.25}; - """) - mad.MADX.load("'test.seq'") - self.assertEqual(mad.MADX.qd.l, 1) - self.assertEqual(mad.MADX.qd.knl.eval(), [0, 0.25]) - seq_file.unlink() - - def test_evaluate_in_madx_environment(self): - with MAD() as mad: - madx_code = """ + """ + ) + mad.MADX.load(f"'{seq_file}'") + assert mad.MADX.qd.l == 1 + assert mad.MADX.qd.knl.eval() == [0, 0.25] + + +def test_evaluate_in_madx_environment(): + with MAD() as mad: + madx_code = """ qd = quadrupole {l=1, knl:={0, 0.25}} """ - mad.evaluate_in_madx_environment(madx_code) - self.assertEqual(mad.MADX.qd.l, 1) - self.assertEqual(mad.MADX.qd.knl.eval(), [0, 0.25]) - - -class TestOps(unittest.TestCase): - def test_matrix(self): - with MAD(py_name="python") as mad: - mad.load("MAD", "matrix") - py_mat = np.arange(1, 101).reshape((10, 10)) - - mad["mat"] = mad.matrix(10).seq(2) + 2 - self.assertTrue(np.all(mad.mat == py_mat + 4)) - - mad["mat"] = mad.matrix(10).seq() / 3 - self.assertTrue(np.allclose(mad.mat, py_mat / 3)) - - mad["mat"] = mad.matrix(10).seq() * 4 - self.assertTrue(np.all(mad.mat == py_mat * 4)) - - mad["mat"] = mad.matrix(10).seq() ** 3 - self.assertTrue(np.all(mad.mat == np.linalg.matrix_power(py_mat, 3))) - - mad["mat"] = mad.matrix(10).seq() + 2 / 3 * 4**3 # bidmas - self.assertTrue(np.all(mad.mat == py_mat + 2 / 3 * 4**3)) - - # conversions - self.assertTrue(np.all(np.array(mad.MAD.matrix(10).seq()) == np.arange(1, 101))) - self.assertTrue(np.all(list(mad.MAD.matrix(10).seq()) == np.arange(1, 101))) - self.assertTrue(np.all(mad.MAD.matrix(10).seq().eval() == py_mat)) - self.assertEqual(np.sin(1), mad.math["sin"](1).eval()) - self.assertAlmostEqual(np.cos(0.5), mad.math["cos"](0.5).eval(), None, None, 4e-16) - - # temp vars - res = ( - ( - (mad.matrix(3).seq().emul(2) * mad.matrix(3).seq(3) + 3) * 2 - + mad.matrix(3).seq(2) - ) - - mad.matrix(3).seq(4) - ).eval() - np_mat = np.arange(9).reshape((3, 3)) + 1 - exp = ((np.matmul((np_mat * 2), (np_mat + 3)) + 3) * 2 + (np_mat + 2)) - (np_mat + 4) - self.assertTrue(np.all(exp == res)) - - -class TestArgsAndKwargs(unittest.TestCase): - def test_args(self): - with MAD() as mad: - mad.load("MAD", "matrix", "cmatrix") - mad["m1"] = mad.matrix(3).seq() - mad["m2"] = mad.matrix(3).eye(2).mul(mad.m1) - mad["m3"] = mad.matrix(3).seq().emul(mad.m1) - mad["cm1"] = mad.cmatrix(3).seq(1j).map([1, 5, 9], "\\mn, n -> mn - 1i") - self.assertTrue(np.all(mad.m2 == mad.m1 * 2)) - self.assertTrue(np.all(mad.m3 == (mad.m1 * mad.m1))) - self.assertTrue(np.all(mad.cm1 == mad.m1 + 1j - np.eye(3) * 1j)) - # Add bool - - def test_kwargs(self): - with MAD() as mad: - mad.load("element", "sextupole") - mad["m1"] = mad.MAD.matrix(3).seq() - sd = mad.sextupole( - knl=[0, 0.25j, 1 + 1j], - l=1, - alist=[1, 2, 3, 5], - abool=True, - opposite=False, - mat=mad.m1, - ) - self.assertEqual(sd.knl.eval(), [0, 0.25j, 1 + 1j]) - self.assertEqual(sd.l, 1) - self.assertEqual(sd.alist.eval(), [1, 2, 3, 5]) - self.assertEqual(sd.abool, True) - self.assertEqual(sd.opposite, False) - self.assertTrue(np.all(sd.mat == np.arange(9).reshape((3, 3)) + 1)) - - -class TestDir(unittest.TestCase): - def test_dir(self): - with MAD(py_name="python") as mad: - mad.load("MAD", "gfunc", "element", "object") - mad.load("element", "quadrupole") - mad.load("gfunc", "functor") - obj_dir = dir(mad.object) - mad.send("my_obj = object {a1 = 2, a2 = functor(\\s->s.a1), a3 = \\s->s.a1}") - obj_exp = sorted(["a1", "a2()", "a3"] + obj_dir) - self.assertEqual(dir(mad.my_obj), obj_exp) - self.assertEqual(mad.my_obj.a1, mad.my_obj.a3) - - quad_exp = dir(mad.quadrupole) - self.assertEqual( - dir(mad.quadrupole(knl=[0, 0.3], l=1)), quad_exp - ) # Dir of instance of class should be the same as the class - self.assertEqual( - dir(mad.quadrupole(asd=10, qwe=20)), sorted(quad_exp + ["asd", "qwe"]) - ) # Adding to the instance should change the dir - - def test_dir_on_mad_object(self): - with MAD(py_name="python") as mad: - mad.load("MAD", "object") - mad.send("my_obj = object {a = 1, b = 2, c = 3}") - expected_dir = sorted(["a", "b", "c"] + dir(mad.object)) - self.assertGreater(len(expected_dir), 0) - self.assertEqual(dir(mad.my_obj), expected_dir) - - def test_dir_on_last_object(self): - with MAD() as mad: - mad.load("MAD", "object") - mad.send("last_obj = __mklast__(object {x = 10, y = 20})") - mad["last"] = mad._MAD__get_MadRef("last_obj") - expected_dir = sorted(["x", "y"] + dir(mad.object)) - self.assertGreater(len(expected_dir), 0) - self.assertEqual(dir(mad.last), expected_dir) - - def test_history(self): - with MAD(debug=True, stdout="/dev/null") as mad: - mad.send("a = 1") - mad.send("b = 2") - mad.send("c = a + b") - history = mad.history() - self.assertIn("a = 1", history) - self.assertIn("b = 2", history) - self.assertIn("c = a + b", history) - - -class TestDataFrame(unittest.TestCase): - def gen_data_frame(self, headers, dataframe, force_pandas=False): - mad = MAD() - mad.send(""" + mad.evaluate_in_madx_environment(madx_code) + assert mad.MADX.qd.l == 1 + assert mad.MADX.qd.knl.eval() == [0, 0.25] + + +def test_matrix(): + with MAD(py_name="python") as mad: + mad.load("MAD", "matrix") + py_mat = np.arange(1, 101).reshape((10, 10)) + + mad["mat"] = mad.matrix(10).seq(2) + 2 + assert np.all(mad.mat == py_mat + 4) + + mad["mat"] = mad.matrix(10).seq() / 3 + assert np.allclose(mad.mat, py_mat / 3) + + mad["mat"] = mad.matrix(10).seq() * 4 + assert np.all(mad.mat == py_mat * 4) + + mad["mat"] = mad.matrix(10).seq() ** 3 + assert np.all(mad.mat == np.linalg.matrix_power(py_mat, 3)) + + mad["mat"] = mad.matrix(10).seq() + 2 / 3 * 4**3 + assert np.all(mad.mat == py_mat + 2 / 3 * 4**3) + + assert np.all(np.array(mad.MAD.matrix(10).seq()) == np.arange(1, 101)) + assert np.all(list(mad.MAD.matrix(10).seq()) == np.arange(1, 101)) + assert np.all(mad.MAD.matrix(10).seq().eval() == py_mat) + assert np.sin(1) == mad.math["sin"](1).eval() + assert mad.math["cos"](0.5).eval() == pytest.approx(np.cos(0.5), abs=4e-16) + + res = ( + ((mad.matrix(3).seq().emul(2) * mad.matrix(3).seq(3) + 3) * 2 + mad.matrix(3).seq(2)) + - mad.matrix(3).seq(4) + ).eval() + np_mat = np.arange(9).reshape((3, 3)) + 1 + exp = ((np.matmul((np_mat * 2), (np_mat + 3)) + 3) * 2 + (np_mat + 2)) - (np_mat + 4) + assert np.all(exp == res) + + +def test_args(): + with MAD() as mad: + mad.load("MAD", "matrix", "cmatrix") + mad["m1"] = mad.matrix(3).seq() + mad["m2"] = mad.matrix(3).eye(2).mul(mad.m1) + mad["m3"] = mad.matrix(3).seq().emul(mad.m1) + mad["cm1"] = mad.cmatrix(3).seq(1j).map([1, 5, 9], "\\mn, n -> mn - 1i") + assert np.all(mad.m2 == mad.m1 * 2) + assert np.all(mad.m3 == (mad.m1 * mad.m1)) + assert np.all(mad.cm1 == mad.m1 + 1j - np.eye(3) * 1j) + + +def test_kwargs(): + with MAD() as mad: + mad.load("element", "sextupole") + mad["m1"] = mad.MAD.matrix(3).seq() + sd = mad.sextupole( + knl=[0, 0.25j, 1 + 1j], + l=1, + alist=[1, 2, 3, 5], + abool=True, + opposite=False, + mat=mad.m1, + ) + assert sd.knl.eval() == [0, 0.25j, 1 + 1j] + assert sd.l == 1 + assert sd.alist.eval() == [1, 2, 3, 5] + assert sd.abool + assert not sd.opposite + assert np.all(sd.mat == np.arange(9).reshape((3, 3)) + 1) + + +def test_dir(): + with MAD(py_name="python") as mad: + mad.load("MAD", "gfunc", "element", "object") + mad.load("element", "quadrupole") + mad.load("gfunc", "functor") + obj_dir = dir(mad.object) + mad.send("my_obj = object {a1 = 2, a2 = functor(\\s->s.a1), a3 = \\s->s.a1}") + obj_exp = sorted(["a1", "a2()", "a3"] + obj_dir) + assert dir(mad.my_obj) == obj_exp + assert mad.my_obj.a1 == mad.my_obj.a3 + + quad_exp = dir(mad.quadrupole) + assert dir(mad.quadrupole(knl=[0, 0.3], l=1)) == quad_exp + assert dir(mad.quadrupole(asd=10, qwe=20)) == sorted(quad_exp + ["asd", "qwe"]) + + +def test_dir_on_mad_object(): + with MAD(py_name="python") as mad: + mad.load("MAD", "object") + mad.send("my_obj = object {a = 1, b = 2, c = 3}") + expected_dir = sorted(["a", "b", "c"] + dir(mad.object)) + assert len(expected_dir) > 0 + assert dir(mad.my_obj) == expected_dir + + +def test_dir_on_last_object(): + with MAD() as mad: + mad.load("MAD", "object") + mad.send("last_obj = __mklast__(object {x = 10, y = 20})") + mad["last"] = mad._MAD__get_MadRef("last_obj") + expected_dir = sorted(["x", "y"] + dir(mad.object)) + assert len(expected_dir) > 0 + assert dir(mad.last) == expected_dir + + +def test_history(): + with MAD(debug=True, stdout="/dev/null") as mad: + mad.send("a = 1") + mad.send("b = 2") + mad.send("c = a + b") + history = mad.history() + assert "a = 1" in history + assert "b = 2" in history + assert "c = a + b" in history + + +def _generate_data_frame(headers_attr, dataframe_type, force_pandas=False): + with MAD() as mad: + mad.send( + """ test = mtable{ - {"string"}, "number", "integer", "complex", "boolean", "list", "table", "range",! "generator", + {"string"}, "number", "integer", "complex", "boolean", "list", "table", "range", name = "test", header = {"string", "number", "integer", "complex", "boolean", "list", "table", "range"}, string = "string", @@ -382,103 +404,105 @@ def gen_data_frame(self, headers, dataframe, force_pandas=False): test:addcol("generator", \\ri, m -> m:getcol("number")[ri] + 1i * m:getcol("number")[ri]) test:write("test") - """) + """ + ) df = mad.test.to_df(force_pandas=force_pandas) - self.assertTrue(isinstance(df, dataframe)) - header = getattr(df, headers) - self.assertEqual(header["name"], "test") - self.assertEqual(header["string"], "string") - self.assertEqual(header["number"], 1.234567890) - self.assertEqual(header["integer"], 12345670) - self.assertEqual(header["complex"], 1.3 + 1.2j) - self.assertEqual(header["boolean"], True) - self.assertEqual(header["list"], [1, 2, 3, 4, 5]) - tbl = getattr(df, headers)["table"] - self.assertEqual(tbl[1], 1) - self.assertEqual(tbl[2], 2) - self.assertEqual(tbl["key"], "value") - self.assertTrue(isinstance(tbl, dict)) - - self.assertEqual(df["string"].tolist(), ["a", "b", "c", "d", "e"]) - self.assertEqual(df["number"].tolist(), [1.1, 2.2, 3.3, 4.4, 5.5]) - self.assertEqual(df["integer"].tolist(), [1, 2, 3, 4, 5]) - self.assertEqual(df["complex"].tolist(), [1 + 2j, 2 + 3j, 3 + 4j, 4 + 5j, 5 + 6j]) - self.assertEqual(df["boolean"].tolist(), [True, False, True, False, True]) + assert isinstance(df, dataframe_type) + header = getattr(df, headers_attr) + assert header["name"] == "test" + assert header["string"] == "string" + assert header["number"] == 1.234567890 + assert header["integer"] == 12345670 + assert header["complex"] == 1.3 + 1.2j + assert header["boolean"] + assert header["list"] == [1, 2, 3, 4, 5] + tbl = header["table"] + assert tbl[1] == 1 + assert tbl[2] == 2 + assert tbl["key"] == "value" + assert isinstance(tbl, dict) + + assert df["string"].tolist() == ["a", "b", "c", "d", "e"] + assert df["number"].tolist() == [1.1, 2.2, 3.3, 4.4, 5.5] + assert df["integer"].tolist() == [1, 2, 3, 4, 5] + assert df["complex"].tolist() == [1 + 2j, 2 + 3j, 3 + 4j, 4 + 5j, 5 + 6j] + assert df["boolean"].tolist() == [True, False, True, False, True] lists = [the_list.eval() for the_list in df["list"]] - self.assertEqual(lists, [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]) - tbl = df["table"].tolist() - for i in range(len(tbl)): - lst = tbl[i] - self.assertEqual([lst[0], lst[1]], [i * 3 + 1, i * 3 + 2]) - self.assertEqual(lst[str((i + 1) * 3)], (i + 1) * 3) - self.assertEqual( - df["range"].tolist(), - [range(1, 12), range(2, 13), range(3, 14), range(4, 15), range(5, 16)], - ) + assert lists == [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]] + tables = df["table"].tolist() + for i, table in enumerate(tables): + assert [table[0], table[1]] == [i * 3 + 1, i * 3 + 2] + assert table[str((i + 1) * 3)] == (i + 1) * 3 + assert df["range"].tolist() == [ + range(1, 12), + range(2, 13), + range(3, 14), + range(4, 15), + range(5, 16), + ] + + +def test_generate_tfs_data_frame(): + _generate_data_frame("headers", tfs.TfsDataFrame) - def test_generate_tfs_data_frame(self): - self.gen_data_frame("headers", tfs.TfsDataFrame) - def test_data_frame_without_tfs(self): - sys.modules["tfs"] = None # Remove tfs-pandas - self.gen_data_frame("attrs", pd.DataFrame) - del sys.modules["tfs"] +def test_data_frame_without_tfs(monkeypatch): + monkeypatch.setitem(sys.modules, "tfs", None) + _generate_data_frame("attrs", pd.DataFrame) - def test_force_pandas_data_frame(self): - self.gen_data_frame("attrs", pd.DataFrame, force_pandas=True) - def test_failure(self): - with MAD() as mad: - mad.send(""" +def test_force_pandas_data_frame(): + _generate_data_frame("attrs", pd.DataFrame, force_pandas=True) + + +def test_data_frame_failure(monkeypatch): + with MAD() as mad: + mad.send( + """ test = mtable{"string", "number"} + {"a", 1.1} + {"b", 2.2} - """) - pandas = sys.modules["pandas"] - sys.modules["pandas"] = None - self.assertRaises(ImportError, lambda: mad.test.to_df()) - sys.modules["pandas"] = pandas - df = mad.test.to_df() - self.assertTrue(isinstance(df, tfs.TfsDataFrame)) - self.assertEqual(df["string"].tolist(), ["a", "b"]) - self.assertEqual(df["number"].tolist(), [1.1, 2.2]) - - -class TestEval(unittest.TestCase): - def test_eval(self): - with MAD() as mad: - mad.send("a = 10; b = 20") - result = mad.eval("a + b") - self.assertEqual(result, 30) - - mad.send("c = {1, 2, 3, 4}") - result = mad.eval("c[2]") # Lua indexing starts at 1 - self.assertEqual(result, 2) - - mad.send("d = MAD.matrix(2, 2):seq()") - result = mad.eval("d[2]") # Accessing matrix element - self.assertEqual(result, 2) - - def test_eval_class(self): - with MAD() as mad: - result = mad.math["sqrt"](2) + mad.math["log"](10) - self.assertTrue(isinstance(result, MadLastRef)) - self.assertEqual(result.eval(), math.sqrt(2) + math.log(10)) - - -class TestIteration(unittest.TestCase): - def test_iterate_through_object(self): - with MAD() as mad: - mad.send(""" + """ + ) + monkeypatch.setitem(sys.modules, "pandas", None) + with pytest.raises(ImportError): + mad.test.to_df() + monkeypatch.setitem(sys.modules, "pandas", pd) + df = mad.test.to_df() + assert isinstance(df, tfs.TfsDataFrame) + assert df["string"].tolist() == ["a", "b"] + assert df["number"].tolist() == [1.1, 2.2] + + +def test_eval(): + with MAD() as mad: + mad.send("a = 10; b = 20") + assert mad.eval("a + b") == 30 + + mad.send("c = {1, 2, 3, 4}") + assert mad.eval("c[2]") == 2 + + mad.send("d = MAD.matrix(2, 2):seq()") + assert mad.eval("d[2]") == 2 + + +def test_eval_class(): + with MAD() as mad: + result = mad.math["sqrt"](2) + mad.math["log"](10) + assert isinstance(result, MadLastRef) + assert result.eval() == math.sqrt(2) + math.log(10) + + +def test_iterate_through_object(): + with MAD() as mad: + mad.send( + """ qd = MAD.element.quadrupole "qd" {knl = {0, 0.25}, l = 1.6} my_obj = sequence {qd, qd, qd, qd} - """) - for elem in mad.my_obj: - if elem.name == "qd": - self.assertEqual(elem.l, 1.6) - for i in range(2): - self.assertEqual(elem.knl[i], 0.25 if i == 1 else 0) - else: - self.assertEqual(elem.kind, "marker") - - -if __name__ == "__main__": - unittest.main() + """ + ) + for elem in mad.my_obj: + if elem.name == "qd": + assert elem.l == 1.6 + for i in range(2): + assert elem.knl[i] == (0.25 if i == 1 else 0) + else: + assert elem.kind == "marker"