[
{
"path": ".github/ISSUE_TEMPLATE/bug_report.md",
"content": "---\nname: Bug report\nabout: Create a report to help us improve\ntitle: ''\nlabels: triage\nassignees: odashi\n\n---\n\n## Environment\n\nIf you used latexify on the browser, fill the following items.\n\n* Browser: \n* Frontend: \n\nIf you used latexify in your own environment, fill at least the following items.\nFeel free to add other items if you think they are useful.\n\n* OS: \n* Python: \n* Package manager: \n* Latexify version: \n\n\n## Description\n\nDescribe the details of the issue. Feel free to insert screenshots if they are useful.\n\n\n## Reproduction\n\nDescribe how to reproduce the issue by other people.\n\n\n## Expected behavior\n\nDescribe how latexify should behave in the case above.\n"
},
{
"path": ".github/ISSUE_TEMPLATE/feature_request.md",
"content": "---\nname: Feature request\nabout: Suggest an idea for this project\ntitle: ''\nlabels: feature\nassignees: odashi\n\n---\n\n## Description\n\nIs your feature request related to a problem? Please describe it.\nA clear and concise description is recommended to proceed the discussion efficiently.\n\n\n## Ideas of the solution\n\nIf you have an idea about the solution you'd like, describe details about it.\n\n\n## Alternative ideas\n\nIf you have other ideas that are already considered, describe them as well.\nThese ideas may also help us to make reasonable decisions.\n\n\n## Additional context\n\nAdd any other context or screenshots about the feature request here.\n"
},
{
"path": ".github/pull_request_template.md",
"content": "\n\n# Overview\n\n\n\n# Details\n\n\n\n# References\n\n\n\n# Blocked by\n\n\n"
},
{
"path": ".github/workflows/ci.yml",
"content": "name: Continuous integration\n\non:\n push:\n branches:\n - main\n pull_request:\n branches: [\"**\"]\n\njobs:\n unit-tests:\n runs-on: ubuntu-latest\n strategy:\n matrix:\n python-version: [\"3.9\", \"3.10\", \"3.11\", \"3.12\", \"3.13\"]\n steps:\n - uses: actions/checkout@v3\n - name: Set up Python ${{ matrix.python-version }}\n uses: actions/setup-python@v4\n with:\n python-version: ${{ matrix.python-version }}\n - name: Install dependencies\n run: |\n python -m pip install --upgrade pip\n python -m pip install -e \".[dev]\"\n - name: Test\n run: python -m pytest src\n\n black:\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v3\n - name: Set up Python\n uses: actions/setup-python@v4\n with:\n python-version: \"3.11\"\n - name: Install dependencies\n run: |\n python -m pip install --upgrade pip\n python -m pip install black\n - name: Check\n run: python -m black -v --check src\n\n flake8:\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v3\n - name: Set up Python\n uses: actions/setup-python@v4\n with:\n python-version: \"3.11\"\n - name: Install dependencies\n run: |\n python -m pip install --upgrade pip\n python -m pip install pyproject-flake8\n - name: Check\n run: pflake8 -v src\n isort:\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v3\n - name: Set up Python\n uses: actions/setup-python@v4\n with:\n python-version: \"3.11\"\n - name: Install dependencies\n run: |\n python -m pip install --upgrade pip\n python -m pip install isort\n - name: Check\n run: python -m isort --check src\n mypy:\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v3\n - name: Set up Python\n uses: actions/setup-python@v4\n with:\n python-version: \"3.11\"\n - name: Install dependencies\n run: |\n python -m pip install --upgrade pip\n python -m pip install '.[mypy]'\n - name: Check\n run: python -m mypy src\n"
},
{
"path": ".github/workflows/release.yml",
"content": "name: Release workflow\n\non:\n push:\n tags:\n - \"v[0123456789].*\"\n\njobs:\n release:\n runs-on: ubuntu-latest\n steps:\n - name: checkout\n uses: actions/checkout@v3\n - name: setup python\n uses: actions/setup-python@v2\n with:\n python-version: \"3.10\"\n - name: build\n run: |\n python -m pip install --upgrade build hatch\n python -m hatch version \"${GITHUB_REF_NAME}\"\n python -m build\n - name: publish\n uses: pypa/gh-action-pypi-publish@release/v1\n with:\n password: ${{ secrets.PYPI_API_TOKEN }}\n"
},
{
"path": ".gitignore",
"content": "# Temporary files\n.swp\ntemp\ntmp\n\n# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\ncover/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\n.pybuilder/\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n# For a library or package, you might want to ignore these files since the code is\n# intended to run in multiple environments; otherwise, check them in:\n.python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# PyCharm project settings\n.idea/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n\n# pytype static type analyzer\n.pytype/\n\n# Cython debug symbols\ncython_debug/\n"
},
{
"path": "CODEOWNERS",
"content": "* @odashi\n"
},
{
"path": "CONTRIBUTING.md",
"content": "# How to Contribute\n\nWe'd love to accept your patches and contributions to this project. There are\njust a few small guidelines you need to follow.\n\n## Contributor License Agreement\n\nContributions to this project must be accompanied by a Contributor License\nAgreement (CLA). You (or your employer) retain the copyright to your\ncontribution; this simply gives us permission to use and redistribute your\ncontributions as part of the project. Head over to\n to see your current agreements on file or\nto sign a new one.\n\nYou generally only need to submit a CLA once, so if you've already submitted one\n(even if it was for a different project), you probably don't need to do it\nagain.\n\n## Code reviews\n\nAll submissions, including submissions by project members, require review. We\nuse GitHub pull requests for this purpose. Consult\n[GitHub Help](https://help.github.com/articles/about-pull-requests/) for more\ninformation on using pull requests.\n\n## Community Guidelines\n\nThis project follows\n[Google's Open Source Community Guidelines](https://opensource.google/conduct/).\n\n## Coding style\n\nThis project follows\n[Tensorflow's style](https://www.tensorflow.org/community/contribute/code_style).\n"
},
{
"path": "LICENSE",
"content": " Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [yyyy] [name of copyright owner]\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n"
},
{
"path": "README.md",
"content": "# latexify\n\n[](https://pypi.org/project/latexify-py/)\n[](https://pypi.org/project/latexify-py/)\n[](https://github.com/google/latexify_py/blob/main/LICENSE)\n[](https://pepy.tech/project/latexify-py)\n[](https://github.com/psf/black)\n[](https://pycqa.github.io/isort/)\n\n`latexify` is a Python package to compile a fragment of Python source code to a\ncorresponding $\\LaTeX$ expression:\n\n\n\n`latexify` provides the following functionalities:\n\n* Libraries to compile Python source code or AST to $\\LaTeX$.\n* IPython classes to pretty-print compiled functions.\n\n## FAQs\n\n1. *Which Python versions are supported?*\n\n Syntaxes on **Pythons 3.9 to 3.13** are officially supported, or will be supported.\n\n2. *Which technique is used?*\n\n `latexify` is implemented as a rule-based system on the official `ast` package.\n\n3. *Are \"AI\" techniques adopted?*\n\n `latexify` is based on traditional parsing techniques.\n If the \"AI\" meant some techniques around machine learning, the answer is no.\n\n## Getting started\n\nSee the\n[example notebook](https://github.com/google/latexify_py/blob/main/examples/latexify_examples.ipynb),\nwhich provides several\nuse-cases of this library.\n\nYou can also try the above notebook on\n[Google Colaboratory](https://colab.research.google.com/github/google/latexify_py/blob/main/examples/latexify_examples.ipynb).\n\nSee also the official\n[documentation](https://github.com/google/latexify_py/blob/main/docs/index.md)\nfor more details.\n\n## How to Contribute\n\nTo contribute to this project, please refer\n[CONTRIBUTING.md](https://github.com/google/latexify_py/blob/develop/CONTRIBUTING.md).\n\n## Disclaimer\n\nThis software is currently hosted on , but not officially\nsupported by Google.\n\nIf you have any issues and/or questions about this software, please visit the\n[issue tracker](https://github.com/google/latexify_py/issues)\nor contact the [main maintainer](https://github.com/odashi).\n\n## License\n\nThis software adopts the\n[Apache License 2.0](https://github.com/google/latexify_py/blob/develop/LICENSE).\n"
},
{
"path": "checks.sh",
"content": "#!/bin/bash\nset -eoux pipefail\n\npython -m pytest src -vv\npython -m black --check src\npython -m pflake8 src\npython -m isort --check src\npython -m mypy src\n"
},
{
"path": "docs/getting_started.md",
"content": "# Getting started\n\nThis document describes how to use `latexify` with your Python code.\n\n\n## Installation\n\n`latexify` depends on only Python libraries at this point.\nYou can simply install `latexify` via `pip`:\n\n```shell\n$ pip install latexify-py\n```\n\nNote that you have to install `latexify-py` rather than `latexify`.\n\n\n## Using `latexify` in Jupyter\n\n`latexify.function` decorator function wraps your functions to pretty-print them as\ncorresponding LaTeX formulas.\nJupyter recognizes this wrapper and try to print LaTeX instead of the original function.\n\nThe following snippet:\n\n```python\n@latexify.function\ndef solve(a, b, c):\n return (-b + math.sqrt(b**2 - 4 * a * c)) / (2 * a)\n\nsolve\n```\n\nwill print the following formula to the output:\n\n$$ \\mathrm{solve}(a, b, c) = \\frac{-b + \\sqrt{b^2 - 4ac}}{2a} $$\n\n\nInvoking wrapped functions work transparently as the original function.\n\n```python\nsolve(1, 2, 1)\n```\n\n```\n-1.0\n```\n\nApplying `str` to the wrapped function returns the underlying LaTeX source.\n\n```python\nprint(solve)\n```\n\n```\nf(n) = \\\\frac{-b + \\\\sqrt{b^{2} - 4ac}}{2a}\n```\n\n`latexify.expression` works similarly to `latexify.function`,\nbut it prints the function without its signature:\n```python\n@latexify.expression\ndef solve(a, b, c):\n return (-b + math.sqrt(b**2 - 4 * a * c)) / (2 * a)\n\nsolve\n```\n\n$$ \\frac{-b + \\sqrt{b^2 - 4ac}}{2a} $$\n\n\n## Obtaining LaTeX expression directly\n\nYou can also use `latexify.get_latex`, which takes a function and directly returns the\nLaTeX expression corresponding to the given function.\n\nThe same parameters with `latexify.function` can be applied to `latexify.get_latex` as\nwell.\n\n```python\ndef solve(a, b, c):\n return (-b + math.sqrt(b**2 - 4 * a * c)) / (2 * a)\n\nlatexify.get_latex(solve)\n```\n\n```\nf(n) = \\\\frac{-b + \\\\sqrt{b^{2} - 4ac}}{2a}\n```\n"
},
{
"path": "docs/index.md",
"content": "# `latexify` documentation\n\n## Index\n\n* [Getting started](getting_started.md)\n* [Parameters](parameters.md)\n\n## External resources\n\n* [Examples on Google Colaboratory](https://colab.research.google.com/drive/1MuiawKpVIZ12MWwyYuzZHmbKThdM5wNJ?usp=sharing)\n"
},
{
"path": "docs/parameters.md",
"content": "# `latexify` parameters\n\nThis document describes the list of parameters to control the behavior of `latexify`.\n\n\n## `identifiers: dict[str, str]`\n\nKey-value pair of identifiers to replace.\n\n```python\nidentifiers = {\n \"my_function\": \"f\",\n \"my_inner_function\": \"g\",\n \"my_argument\": \"x\",\n}\n\n@latexify.function(identifiers=identifiers)\ndef my_function(my_argument):\n return my_inner_function(my_argument)\n\nmy_function\n```\n\n$$f(x) = \\mathrm{g}\\left(x\\right)$$\n\n\n## `reduce_assignments: bool`\n\nWhether to compose all variables defined before the `return` statement.\n\nThe current version of `latexify` recognizes only the assignment statements.\nAnalyzing functions with other control flows may raise errors.\n\n```python\n@latexify.function(reduce_assignments=True)\ndef f(a, b, c):\n discriminant = b**2 - 4 * a * c\n numerator = -b + math.sqrt(discriminant)\n denominator = 2 * a\n return numerator / denominator\n\nf\n```\n\n$$f(a, b, c) = \\frac{-b + \\sqrt{b^{2} - 4 a c}}{2 a}$$\n\n\n## `use_math_symbols: bool`\n\nWhether to automatically convert variables with symbol names into LaTeX symbols or not.\n\n```python\n@latexify.function(use_math_symbols=True)\ndef greek(alpha, beta, gamma, Omega):\n return alpha * beta + math.gamma(gamma) + Omega\n\ngreek\n```\n\n$$\\mathrm{greek}({\\alpha}, {\\beta}, {\\gamma}, {\\Omega}) = {\\alpha} {\\beta} + \\Gamma\\left({{\\gamma}}\\right) + {\\Omega}$$\n\n\n## `use_set_symbols: bool`\n\nWhether to use binary operators for set operations or not.\n\n```python\n@latexify.function(use_set_symbols=True)\ndef f(x, y):\n return x & y, x | y, x - y, x ^ y, x < y, x <= y, x > y, x >= y\n\nf\n```\n\n$$f(x, y) = \\left( x \\cap y\\space,\\space x \\cup y\\space,\\space x \\setminus y\\space,\\space x \\mathbin{\\triangle} y\\space,\\space {x \\subset y}\\space,\\space {x \\subseteq y}\\space,\\space {x \\supset y}\\space,\\space {x \\supseteq y}\\right)$$\n\n\n## `use_signature: bool`\n\nWhether to output the function signature or not.\n\nThe default value of this flag depends on the frontend function.\n`True` is used in `latexify.function`, while `False` is used in `latexify.expression`.\n\n```python\n@latexify.function(use_signature=False)\ndef f(a, b, c):\n return (-b + math.sqrt(b**2 - 4 * a * c)) / (2 * a)\n\nf\n```\n\n$$\\frac{-b + \\sqrt{b^{2} - 4 a c}}{2 a}$$\n"
},
{
"path": "examples/latexify_examples.ipynb",
"content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"W5mNJI3Bnl6n\"\n },\n \"source\": [\n \"# `latexify` examples\\n\",\n \"\\n\",\n \"This notebook provides several examples to use `latexify`.\\n\",\n \"\\n\",\n \"See also the\\n\",\n \"[official documentation](https://github.com/google/latexify_py/blob/documentation/docs/index.md)\\n\",\n \"for more details.\\n\",\n \"\\n\",\n \"If you have any questions, please ask it in the\\n\",\n \"[issue tracker](https://github.com/google/latexify_py/issues).\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"fWCVgcRHoLd8\"\n },\n \"source\": [\n \"## Install `latexify`\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\"\n },\n \"id\": \"4IPGyu2dFH6T\",\n \"outputId\": \"471cab8d-3069-4a27-f3ff-67ba177ec58d\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Collecting latexify-py\\n\",\n \" Downloading latexify_py-0.4.2-py3-none-any.whl (38 kB)\\n\",\n \"Collecting dill>=0.3.2 (from latexify-py)\\n\",\n \" Downloading dill-0.3.7-py3-none-any.whl (115 kB)\\n\",\n \"\\u001b[2K \\u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\\u001b[0m \\u001b[32m115.3/115.3 kB\\u001b[0m \\u001b[31m5.7 MB/s\\u001b[0m eta \\u001b[36m0:00:00\\u001b[0m\\n\",\n \"\\u001b[?25hInstalling collected packages: dill, latexify-py\\n\",\n \"Successfully installed dill-0.3.7 latexify-py-0.4.2\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Restart the runtime before running the examples below.\\n\",\n \"%pip install latexify-py\\n\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"-Mzq4_dNoSmc\"\n },\n \"source\": [\n \"## Import `latexify` into your code\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\",\n \"height\": 35\n },\n \"id\": \"hViDMhyMFNCO\",\n \"outputId\": \"b46edb25-5952-4cff-da1e-d65e7e3caad0\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"application/vnd.google.colaboratory.intrinsic+json\": {\n \"type\": \"string\"\n },\n \"text/plain\": [\n \"'0.4.2'\"\n ]\n },\n \"execution_count\": 2,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"import math # Optional\\n\",\n \"import numpy as np # Optional\\n\",\n \"import latexify\\n\",\n \"\\n\",\n \"latexify.__version__\\n\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"4QJ6I2s7odX1\"\n },\n \"source\": [\n \"## Examples\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\"\n },\n \"id\": \"NvbEYSwXFaeE\",\n \"outputId\": \"5d0ca2a4-a285-4053-9cc4-3776746443be\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"-1.0\\n\",\n \"\\\\mathrm{solve}(a, b, c) = \\\\frac{-b + \\\\sqrt{ b^{2} - 4 a c }}{2 a}\\n\"\n ]\n }\n ],\n \"source\": [\n \"@latexify.function\\n\",\n \"def solve(a, b, c):\\n\",\n \" return (-b + math.sqrt(b**2 - 4*a*c)) / (2*a)\\n\",\n \"\\n\",\n \"print(solve(1, 4, 3)) # Invoking the function works as expected.\\n\",\n \"print(solve) # Printing the function shows the underlying LaTeX source.\\n\",\n \"solve # Displays the expression.\\n\",\n \"\\n\",\n \"# Writes the underlying LaTeX source into a file.\\n\",\n \"with open(\\\"compiled.tex\\\", \\\"w\\\") as fp:\\n\",\n \" print(solve, file=fp)\\n\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\",\n \"height\": 56\n },\n \"id\": \"wS7BhtPgjSak\",\n \"outputId\": \"76a8547c-e6b5-458d-aeb2-f9df2f35f7c7\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/latex\": [\n \"$$ \\\\displaystyle \\\\frac{-b + \\\\sqrt{ b^{2} - 4 a c }}{2 a} $$\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"execution_count\": 4,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# latexify.expression works similarly, but does not output the signature.\\n\",\n \"@latexify.expression\\n\",\n \"def solve(a, b, c):\\n\",\n \" return (-b + math.sqrt(b**2 - 4*a*c)) / (2*a)\\n\",\n \"\\n\",\n \"solve\\n\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\",\n \"height\": 35\n },\n \"id\": \"G73dnoqqjg4A\",\n \"outputId\": \"b9f53cf8-4a34-452c-8d9b-946ddd0998df\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"application/vnd.google.colaboratory.intrinsic+json\": {\n \"type\": \"string\"\n },\n \"text/plain\": [\n \"'\\\\\\\\mathrm{solve}(a, b, c) = \\\\\\\\frac{-b + \\\\\\\\sqrt{ b^{2} - 4 a c }}{2 a}'\"\n ]\n },\n \"execution_count\": 5,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# latexify.get_latex obtains the underlying LaTeX expression directly.\\n\",\n \"def solve(a, b, c):\\n\",\n \" return (-b + math.sqrt(b**2 - 4*a*c)) / (2*a)\\n\",\n \"\\n\",\n \"latexify.get_latex(solve)\\n\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\",\n \"height\": 58\n },\n \"id\": \"8bYSWIngGF8E\",\n \"outputId\": \"669e070d-2718-49cb-a2fe-0defe0286b27\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/latex\": [\n \"$$ \\\\displaystyle \\\\mathrm{sinc}(x) = \\\\left\\\\{ \\\\begin{array}{ll} 1, & \\\\mathrm{if} \\\\ x = 0 \\\\\\\\ \\\\frac{\\\\sin x}{x}, & \\\\mathrm{otherwise} \\\\end{array} \\\\right. $$\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"execution_count\": 6,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"@latexify.function\\n\",\n \"def sinc(x):\\n\",\n \" if x == 0:\\n\",\n \" return 1\\n\",\n \" else:\\n\",\n \" return math.sin(x) / x\\n\",\n \"\\n\",\n \"sinc\\n\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\",\n \"height\": 78\n },\n \"id\": \"h1i4BjdgHjxl\",\n \"outputId\": \"e448ff37-4753-4090-b2b1-1ef21b279b34\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/latex\": [\n \"$$ \\\\displaystyle \\\\mathrm{fib}(x) = \\\\left\\\\{ \\\\begin{array}{ll} 0, & \\\\mathrm{if} \\\\ x = 0 \\\\\\\\ 1, & \\\\mathrm{if} \\\\ x = 1 \\\\\\\\ \\\\mathrm{fib} \\\\mathopen{}\\\\left( x - 1 \\\\mathclose{}\\\\right) + \\\\mathrm{fib} \\\\mathopen{}\\\\left( x - 2 \\\\mathclose{}\\\\right), & \\\\mathrm{otherwise} \\\\end{array} \\\\right. $$\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"execution_count\": 7,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# Elif or nested else-if are unrolled.\\n\",\n \"@latexify.function\\n\",\n \"def fib(x):\\n\",\n \" if x == 0:\\n\",\n \" return 0\\n\",\n \" elif x == 1:\\n\",\n \" return 1\\n\",\n \" else:\\n\",\n \" return fib(x-1) + fib(x-2)\\n\",\n \"\\n\",\n \"fib\\n\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\",\n \"height\": 39\n },\n \"id\": \"-JhJMAXM7j-X\",\n \"outputId\": \"a47dcd59-2ff9-4aa1-935d-7c789b39057e\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/latex\": [\n \"$$ \\\\displaystyle \\\\mathrm{greek}(\\\\alpha, \\\\beta, \\\\gamma, \\\\Omega) = \\\\alpha \\\\beta + \\\\Gamma \\\\mathopen{}\\\\left( \\\\gamma \\\\mathclose{}\\\\right) + \\\\Omega $$\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"execution_count\": 8,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# Some math symbols are converted automatically.\\n\",\n \"@latexify.function(use_math_symbols=True)\\n\",\n \"def greek(alpha, beta, gamma, Omega):\\n\",\n \" return alpha * beta + math.gamma(gamma) + Omega\\n\",\n \"\\n\",\n \"greek\\n\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\",\n \"height\": 39\n },\n \"id\": \"ySyNPS0y4tzu\",\n \"outputId\": \"2d95b5ce-a9b8-42b1-eb55-dc8bd0097d69\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/latex\": [\n \"$$ \\\\displaystyle f(x) = g \\\\mathopen{}\\\\left( x \\\\mathclose{}\\\\right) $$\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"execution_count\": 9,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# Function names, arguments, variables can be replaced.\\n\",\n \"identifiers = {\\n\",\n \" \\\"my_function\\\": \\\"f\\\",\\n\",\n \" \\\"my_inner_function\\\": \\\"g\\\",\\n\",\n \" \\\"my_argument\\\": \\\"x\\\",\\n\",\n \"}\\n\",\n \"\\n\",\n \"@latexify.function(identifiers=identifiers)\\n\",\n \"def my_function(my_argument):\\n\",\n \" return my_inner_function(my_argument)\\n\",\n \"\\n\",\n \"my_function\\n\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\",\n \"height\": 56\n },\n \"id\": \"TyacQaDM4Ei7\",\n \"outputId\": \"8e971bbd-2c74-45d2-d0fa-7f46569b10a6\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/latex\": [\n \"$$ \\\\displaystyle f(a, b, c) = \\\\frac{-b + \\\\sqrt{ b^{2} - 4 a c }}{2 a} $$\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"execution_count\": 10,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# Assignments can be reduced into one expression.\\n\",\n \"@latexify.function(reduce_assignments=True)\\n\",\n \"def f(a, b, c):\\n\",\n \" discriminant = b**2 - 4 * a * c\\n\",\n \" numerator = -b + math.sqrt(discriminant)\\n\",\n \" denominator = 2 * a\\n\",\n \" return numerator / denominator\\n\",\n \"\\n\",\n \"f\\n\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 11,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\",\n \"height\": 78\n },\n \"id\": \"oD8MFS2WE-2U\",\n \"outputId\": \"f9fad1bd-b7eb-41cc-8743-ec0d80cca8bc\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/latex\": [\n \"$$ \\\\displaystyle \\\\mathrm{transform}(x, y, a, b, \\\\theta, s, t) = \\\\begin{bmatrix} 1 & 0 & s \\\\\\\\ 0 & 1 & t \\\\\\\\ 0 & 0 & 1 \\\\end{bmatrix} \\\\cdot \\\\begin{bmatrix} \\\\cos \\\\theta & -\\\\sin \\\\theta & 0 \\\\\\\\ \\\\sin \\\\theta & \\\\cos \\\\theta & 0 \\\\\\\\ 0 & 0 & 1 \\\\end{bmatrix} \\\\cdot \\\\begin{bmatrix} a & 0 & 0 \\\\\\\\ 0 & b & 0 \\\\\\\\ 0 & 0 & 1 \\\\end{bmatrix} \\\\cdot \\\\begin{bmatrix} x \\\\\\\\ y \\\\\\\\ 1 \\\\end{bmatrix} $$\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"execution_count\": 11,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# Matrix support.\\n\",\n \"@latexify.function(reduce_assignments=True, use_math_symbols=True)\\n\",\n \"def transform(x, y, a, b, theta, s, t):\\n\",\n \" cos_t = math.cos(theta)\\n\",\n \" sin_t = math.sin(theta)\\n\",\n \" scale = np.array([[a, 0, 0], [0, b, 0], [0, 0, 1]])\\n\",\n \" rotate = np.array([[cos_t, -sin_t, 0], [sin_t, cos_t, 0], [0, 0, 1]])\\n\",\n \" move = np.array([[1, 0, s], [0, 1, t], [0, 0, 1]])\\n\",\n \" return move @ rotate @ scale @ np.array([[x], [y], [1]])\\n\",\n \"\\n\",\n \"transform\\n\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\",\n \"height\": 241\n },\n \"id\": \"81OlPVWyGfWN\",\n \"outputId\": \"48660400-a812-41e2-91ea-23e49ea20c7f\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/latex\": [\n \"$ \\\\begin{array}{l} \\\\mathbf{function} \\\\ \\\\mathrm{fib}(x) \\\\\\\\ \\\\hspace{1em} \\\\mathbf{if} \\\\ x = 0 \\\\\\\\ \\\\hspace{2em} \\\\mathbf{return} \\\\ 0 \\\\\\\\ \\\\hspace{1em} \\\\mathbf{else} \\\\\\\\ \\\\hspace{2em} \\\\mathbf{if} \\\\ x = 1 \\\\\\\\ \\\\hspace{3em} \\\\mathbf{return} \\\\ 1 \\\\\\\\ \\\\hspace{2em} \\\\mathbf{else} \\\\\\\\ \\\\hspace{3em} \\\\mathbf{return} \\\\ \\\\mathrm{fib} \\\\mathopen{}\\\\left( x - 1 \\\\mathclose{}\\\\right) + \\\\mathrm{fib} \\\\mathopen{}\\\\left( x - 2 \\\\mathclose{}\\\\right) \\\\\\\\ \\\\hspace{2em} \\\\mathbf{end \\\\ if} \\\\\\\\ \\\\hspace{1em} \\\\mathbf{end \\\\ if} \\\\\\\\ \\\\mathbf{end \\\\ function} \\\\end{array} $\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"execution_count\": 12,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# latexify.algorithmic generates an algorithmic environment instead of an equation.\\n\",\n \"@latexify.algorithmic\\n\",\n \"def fib(x):\\n\",\n \" if x == 0:\\n\",\n \" return 0\\n\",\n \" elif x == 1:\\n\",\n \" return 1\\n\",\n \" else:\\n\",\n \" return fib(x-1) + fib(x-2)\\n\",\n \"\\n\",\n \"fib\\n\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 13,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\",\n \"height\": 261\n },\n \"id\": \"kbw_1txkGfnX\",\n \"outputId\": \"fdc58207-1c06-4d88-e249-b0b011bd98c0\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/latex\": [\n \"$ \\\\begin{array}{l} \\\\mathbf{function} \\\\ \\\\mathrm{collatz}(x) \\\\\\\\ \\\\hspace{1em} n \\\\gets 0 \\\\\\\\ \\\\hspace{1em} \\\\mathbf{while} \\\\ x > 1 \\\\\\\\ \\\\hspace{2em} n \\\\gets n + 1 \\\\\\\\ \\\\hspace{2em} \\\\mathbf{if} \\\\ x \\\\mathbin{\\\\%} 2 = 0 \\\\\\\\ \\\\hspace{3em} x \\\\gets \\\\left\\\\lfloor\\\\frac{x}{2}\\\\right\\\\rfloor \\\\\\\\ \\\\hspace{2em} \\\\mathbf{else} \\\\\\\\ \\\\hspace{3em} x \\\\gets 3 x + 1 \\\\\\\\ \\\\hspace{2em} \\\\mathbf{end \\\\ if} \\\\\\\\ \\\\hspace{1em} \\\\mathbf{end \\\\ while} \\\\\\\\ \\\\hspace{1em} \\\\mathbf{return} \\\\ n \\\\\\\\ \\\\mathbf{end \\\\ function} \\\\end{array} $\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"execution_count\": 13,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# Another example: latexify.algorithmic supports usual control flows.\\n\",\n \"@latexify.algorithmic\\n\",\n \"def collatz(x):\\n\",\n \" n = 0\\n\",\n \" while x > 1:\\n\",\n \" n = n + 1\\n\",\n \" if x % 2 == 0:\\n\",\n \" x = x // 2\\n\",\n \" else:\\n\",\n \" x = 3 * x + 1\\n\",\n \" return n\\n\",\n \"\\n\",\n \"collatz\\n\"\n ]\n }\n ],\n \"metadata\": {\n \"colab\": {\n \"provenance\": []\n },\n \"kernelspec\": {\n \"display_name\": \"Python 3\",\n \"name\": \"python3\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 0\n}\n"
},
{
"path": "pyproject.toml",
"content": "[build-system]\nrequires = [\n \"hatchling\",\n]\nbuild-backend = \"hatchling.build\"\n\n[project]\nname = \"latexify-py\"\ndescription = \"Generates LaTeX math description from Python functions.\"\nreadme = \"README.md\"\nrequires-python = \">=3.9, <3.14\"\nlicense = {text = \"Apache Software License 2.0\"}\nauthors = [\n {name = \"Yusuke Oda\", email = \"odashi@inspiredco.ai\"}\n]\nkeywords = [\n \"equation\",\n \"latex\",\n \"math\",\n \"mathematics\",\n \"tex\",\n]\nclassifiers = [\n \"Framework :: IPython\",\n \"Framework :: Jupyter\",\n \"License :: OSI Approved :: Apache Software License\",\n \"Programming Language :: Python :: 3.9\",\n \"Programming Language :: Python :: 3.10\",\n \"Programming Language :: Python :: 3.11\",\n \"Programming Language :: Python :: 3.12\",\n \"Programming Language :: Python :: 3.13\",\n \"Topic :: Scientific/Engineering :: Mathematics\",\n \"Topic :: Software Development :: Code Generators\",\n \"Topic :: Text Processing :: Markup :: LaTeX\",\n]\ndependencies = [\n \"dill>=0.3.2\",\n]\ndynamic = [\n \"version\"\n]\n\n[project.optional-dependencies]\ndev = [\n \"build>=0.8\",\n \"black>=24.3\",\n \"flake8>=6.0\",\n \"isort>=5.10\",\n \"mypy>=1.9\",\n \"notebook>=6.5.1\",\n \"pyproject-flake8>=6.0\",\n \"pytest>=7.1\",\n \"twine>=4.0\",\n]\nmypy = [\n \"mypy>=1.9\",\n \"pytest>=7.1\",\n]\n\n[project.urls]\nHomepage = \"https://github.com/google/latexify_py\"\n\"Bug Tracker\" = \"https://github.com/google/latexify_py/issues\"\n\n[tool.hatch.build]\ninclude = [\n \"*.py\",\n]\nexclude = [\n \"*_test.py\",\n]\nonly-packages = true\n\n[tool.hatch.build.targets.wheel]\npackages = [\"src/latexify\"]\n\n[tool.hatch.version]\npath = \"src/latexify/_version.py\"\n\n[tool.flake8]\nmax-line-length = 88\nextend-ignore = \"E203\"\n\n[tool.isort]\nprofile = \"black\"\n"
},
{
"path": "src/integration_tests/__init__.py",
"content": "\"\"\"Package integration_tests.\"\"\"\n\nimport pytest\n\npytest.register_assert_rewrite(\"integration_tests.utils\")\n"
},
{
"path": "src/integration_tests/algorithmic_style_test.py",
"content": "\"\"\"End-to-end test cases of algorithmic style.\"\"\"\n\nfrom __future__ import annotations\n\nimport textwrap\n\nfrom integration_tests import integration_utils\n\n\ndef test_factorial() -> None:\n def fact(n):\n if n == 0:\n return 1\n else:\n return n * fact(n - 1)\n\n latex = textwrap.dedent(\n r\"\"\"\n \\begin{algorithmic}\n \\Function{fact}{$n$}\n \\If{$n = 0$}\n \\State \\Return $1$\n \\Else\n \\State \\Return $n \\cdot \\mathrm{fact} \\mathopen{}\\left( n - 1 \\mathclose{}\\right)$\n \\EndIf\n \\EndFunction\n \\end{algorithmic}\n \"\"\" # noqa: E501\n ).strip()\n ipython_latex = (\n r\"\\begin{array}{l}\"\n r\" \\mathbf{function} \\ \\mathrm{fact}(n) \\\\\"\n r\" \\hspace{1em} \\mathbf{if} \\ n = 0 \\\\\"\n r\" \\hspace{2em} \\mathbf{return} \\ 1 \\\\\"\n r\" \\hspace{1em} \\mathbf{else} \\\\\"\n r\" \\hspace{2em}\"\n r\" \\mathbf{return} \\ n \\cdot\"\n r\" \\mathrm{fact} \\mathopen{}\\left( n - 1 \\mathclose{}\\right) \\\\\"\n r\" \\hspace{1em} \\mathbf{end \\ if} \\\\\"\n r\" \\mathbf{end \\ function}\"\n r\" \\end{array}\"\n )\n integration_utils.check_algorithm(fact, latex, ipython_latex)\n\n\ndef test_collatz() -> None:\n def collatz(n):\n iterations = 0\n while n > 1:\n if n % 2 == 0:\n n = n // 2\n else:\n n = 3 * n + 1\n iterations = iterations + 1\n return iterations\n\n latex = textwrap.dedent(\n r\"\"\"\n \\begin{algorithmic}\n \\Function{collatz}{$n$}\n \\State $\\mathrm{iterations} \\gets 0$\n \\While{$n > 1$}\n \\If{$n \\mathbin{\\%} 2 = 0$}\n \\State $n \\gets \\left\\lfloor\\frac{n}{2}\\right\\rfloor$\n \\Else\n \\State $n \\gets 3 n + 1$\n \\EndIf\n \\State $\\mathrm{iterations} \\gets \\mathrm{iterations} + 1$\n \\EndWhile\n \\State \\Return $\\mathrm{iterations}$\n \\EndFunction\n \\end{algorithmic}\n \"\"\"\n ).strip()\n ipython_latex = (\n r\"\\begin{array}{l}\"\n r\" \\mathbf{function} \\ \\mathrm{collatz}(n) \\\\\"\n r\" \\hspace{1em} \\mathrm{iterations} \\gets 0 \\\\\"\n r\" \\hspace{1em} \\mathbf{while} \\ n > 1 \\\\\"\n r\" \\hspace{2em} \\mathbf{if} \\ n \\mathbin{\\%} 2 = 0 \\\\\"\n r\" \\hspace{3em} n \\gets \\left\\lfloor\\frac{n}{2}\\right\\rfloor \\\\\"\n r\" \\hspace{2em} \\mathbf{else} \\\\\"\n r\" \\hspace{3em} n \\gets 3 n + 1 \\\\\"\n r\" \\hspace{2em} \\mathbf{end \\ if} \\\\\"\n r\" \\hspace{2em}\"\n r\" \\mathrm{iterations} \\gets \\mathrm{iterations} + 1 \\\\\"\n r\" \\hspace{1em} \\mathbf{end \\ while} \\\\\"\n r\" \\hspace{1em} \\mathbf{return} \\ \\mathrm{iterations} \\\\\"\n r\" \\mathbf{end \\ function}\"\n r\" \\end{array}\"\n )\n integration_utils.check_algorithm(collatz, latex, ipython_latex)\n"
},
{
"path": "src/integration_tests/function_expansion_test.py",
"content": "\"\"\"End-to-end test cases of function expansion.\"\"\"\n\nfrom __future__ import annotations\n\nimport math\n\nfrom integration_tests import integration_utils\n\n\ndef test_atan2() -> None:\n def solve(x, y):\n return math.atan2(y, x)\n\n latex = (\n r\"\\mathrm{solve}(x, y) =\"\n r\" \\arctan \\mathopen{}\\left( \\frac{y}{x} \\mathclose{}\\right)\"\n )\n integration_utils.check_function(solve, latex, expand_functions={\"atan2\"})\n\n\ndef test_atan2_nested() -> None:\n def solve(x, y):\n return math.atan2(math.exp(y), math.exp(x))\n\n latex = (\n r\"\\mathrm{solve}(x, y) =\"\n r\" \\arctan \\mathopen{}\\left( \\frac{e^{y}}{e^{x}} \\mathclose{}\\right)\"\n )\n integration_utils.check_function(solve, latex, expand_functions={\"atan2\", \"exp\"})\n\n\ndef test_exp() -> None:\n def solve(x):\n return math.exp(x)\n\n latex = r\"\\mathrm{solve}(x) = e^{x}\"\n integration_utils.check_function(solve, latex, expand_functions={\"exp\"})\n\n\ndef test_exp_nested() -> None:\n def solve(x):\n return math.exp(math.exp(x))\n\n latex = r\"\\mathrm{solve}(x) = e^{e^{x}}\"\n integration_utils.check_function(solve, latex, expand_functions={\"exp\"})\n\n\ndef test_exp2() -> None:\n def solve(x):\n return math.exp2(x)\n\n latex = r\"\\mathrm{solve}(x) = 2^{x}\"\n integration_utils.check_function(solve, latex, expand_functions={\"exp2\"})\n\n\ndef test_exp2_nested() -> None:\n def solve(x):\n return math.exp2(math.exp2(x))\n\n latex = r\"\\mathrm{solve}(x) = 2^{2^{x}}\"\n integration_utils.check_function(solve, latex, expand_functions={\"exp2\"})\n\n\ndef test_expm1() -> None:\n def solve(x):\n return math.expm1(x)\n\n latex = r\"\\mathrm{solve}(x) = \\exp x - 1\"\n integration_utils.check_function(solve, latex, expand_functions={\"expm1\"})\n\n\ndef test_expm1_nested() -> None:\n def solve(x, y, z):\n return math.expm1(math.pow(y, z))\n\n latex = r\"\\mathrm{solve}(x, y, z) = e^{y^{z}} - 1\"\n integration_utils.check_function(\n solve, latex, expand_functions={\"expm1\", \"exp\", \"pow\"}\n )\n\n\ndef test_hypot_without_attribute() -> None:\n from math import hypot\n\n def solve(x, y, z):\n return hypot(x, y, z)\n\n latex = r\"\\mathrm{solve}(x, y, z) = \\sqrt{ x^{2} + y^{2} + z^{2} }\"\n integration_utils.check_function(solve, latex, expand_functions={\"hypot\"})\n\n\ndef test_hypot() -> None:\n def solve(x, y, z):\n return math.hypot(x, y, z)\n\n latex = r\"\\mathrm{solve}(x, y, z) = \\sqrt{ x^{2} + y^{2} + z^{2} }\"\n integration_utils.check_function(solve, latex, expand_functions={\"hypot\"})\n\n\ndef test_hypot_nested() -> None:\n def solve(a, b, x, y):\n return math.hypot(math.hypot(a, b), x, y)\n\n latex = (\n r\"\\mathrm{solve}(a, b, x, y) =\"\n r\" \\sqrt{ \\sqrt{ a^{2} + b^{2} }^{2} + x^{2} + y^{2} }\"\n )\n integration_utils.check_function(solve, latex, expand_functions={\"hypot\"})\n\n\ndef test_log1p() -> None:\n def solve(x):\n return math.log1p(x)\n\n latex = r\"\\mathrm{solve}(x) = \\log \\mathopen{}\\left( 1 + x \\mathclose{}\\right)\"\n integration_utils.check_function(solve, latex, expand_functions={\"log1p\"})\n\n\ndef test_log1p_nested() -> None:\n def solve(x):\n return math.log1p(math.exp(x))\n\n latex = r\"\\mathrm{solve}(x) = \\log \\mathopen{}\\left( 1 + e^{x} \\mathclose{}\\right)\"\n integration_utils.check_function(solve, latex, expand_functions={\"log1p\", \"exp\"})\n\n\ndef test_pow_nested() -> None:\n def solve(w, x, y, z):\n return math.pow(math.pow(w, x), math.pow(y, z))\n\n latex = (\n r\"\\mathrm{solve}(w, x, y, z) = \"\n r\"\\mathopen{}\\left( w^{x} \\mathclose{}\\right)^{y^{z}}\"\n )\n integration_utils.check_function(solve, latex, expand_functions={\"pow\"})\n\n\ndef test_pow() -> None:\n def solve(x, y):\n return math.pow(x, y)\n\n latex = r\"\\mathrm{solve}(x, y) = x^{y}\"\n integration_utils.check_function(solve, latex, expand_functions={\"pow\"})\n"
},
{
"path": "src/integration_tests/integration_utils.py",
"content": "\"\"\"Utilities for integration tests.\"\"\"\n\nfrom __future__ import annotations\n\nfrom typing import Any, Callable\n\nfrom latexify import frontend\n\n\ndef check_function(\n fn: Callable[..., Any],\n latex: str,\n **kwargs,\n) -> None:\n \"\"\"Helper to check if the obtained function has the expected LaTeX form.\n\n Args:\n fn: Function to check.\n latex: LaTeX form of `fn`.\n **kwargs: Arguments passed to `frontend.function`.\n \"\"\"\n # Checks the syntax:\n # @function\n # def fn(...):\n # ...\n if not kwargs:\n latexified = frontend.function(fn)\n assert str(latexified) == latex\n assert latexified._repr_latex_() == rf\"$$ \\displaystyle {latex} $$\"\n\n # Checks the syntax:\n # @function(**kwargs)\n # def fn(...):\n # ...\n latexified = frontend.function(**kwargs)(fn)\n assert str(latexified) == latex\n assert latexified._repr_latex_() == rf\"$$ \\displaystyle {latex} $$\"\n\n # Checks the syntax:\n # def fn(...):\n # ...\n # latexified = function(fn, **kwargs)\n latexified = frontend.function(fn, **kwargs)\n assert str(latexified) == latex\n assert latexified._repr_latex_() == rf\"$$ \\displaystyle {latex} $$\"\n\n\ndef check_algorithm(\n fn: Callable[..., Any],\n latex: str,\n ipython_latex: str,\n **kwargs,\n) -> None:\n \"\"\"Helper to check if the obtained function has the expected LaTeX form.\n\n Args:\n fn: Function to check.\n latex: LaTeX form of `fn`.\n ipython_latex: IPython LaTeX form of `fn`\n **kwargs: Arguments passed to `frontend.get_latex`.\n \"\"\"\n # Checks the syntax:\n # @algorithmic\n # def fn(...):\n # ...\n if not kwargs:\n latexified = frontend.algorithmic(fn)\n assert str(latexified) == latex\n assert latexified._repr_latex_() == f\"$ {ipython_latex} $\"\n\n # Checks the syntax:\n # @algorithmic(**kwargs)\n # def fn(...):\n # ...\n latexified = frontend.algorithmic(**kwargs)(fn)\n assert str(latexified) == latex\n assert latexified._repr_latex_() == f\"$ {ipython_latex} $\"\n\n # Checks the syntax:\n # def fn(...):\n # ...\n # latexified = algorithmic(fn, **kwargs)\n latexified = frontend.algorithmic(fn, **kwargs)\n assert str(latexified) == latex\n assert latexified._repr_latex_() == f\"$ {ipython_latex} $\"\n"
},
{
"path": "src/integration_tests/regression_test.py",
"content": "\"\"\"End-to-end test cases of function.\"\"\"\n\nfrom __future__ import annotations\n\nimport math\n\nfrom integration_tests import integration_utils\n\n\ndef test_quadratic_solution() -> None:\n def solve(a, b, c):\n return (-b + math.sqrt(b**2 - 4 * a * c)) / (2 * a)\n\n latex = r\"\\mathrm{solve}(a, b, c) =\" r\" \\frac{-b + \\sqrt{ b^{2} - 4 a c }}{2 a}\"\n integration_utils.check_function(solve, latex)\n\n\ndef test_sinc() -> None:\n def sinc(x):\n if x == 0:\n return 1\n else:\n return math.sin(x) / x\n\n latex = (\n r\"\\mathrm{sinc}(x) =\"\n r\" \\left\\{ \\begin{array}{ll}\"\n r\" 1, & \\mathrm{if} \\ x = 0 \\\\\"\n r\" \\frac{\\sin x}{x}, & \\mathrm{otherwise}\"\n r\" \\end{array} \\right.\"\n )\n integration_utils.check_function(sinc, latex)\n\n\ndef test_x_times_beta() -> None:\n def xtimesbeta(x, beta):\n return x * beta\n\n latex_without_symbols = (\n r\"\\mathrm{xtimesbeta}(x, \\mathrm{beta}) = x \\cdot \\mathrm{beta}\"\n )\n integration_utils.check_function(xtimesbeta, latex_without_symbols)\n integration_utils.check_function(\n xtimesbeta, latex_without_symbols, use_math_symbols=False\n )\n\n latex_with_symbols = r\"\\mathrm{xtimesbeta}(x, \\beta) = x \\beta\"\n integration_utils.check_function(\n xtimesbeta, latex_with_symbols, use_math_symbols=True\n )\n\n\ndef test_sum_with_limit_1arg() -> None:\n def sum_with_limit(n):\n return sum(i**2 for i in range(n))\n\n latex = (\n r\"\\mathrm{sum\\_with\\_limit}(n) = \\sum_{i = 0}^{n - 1}\"\n r\" \\mathopen{}\\left({i^{2}}\\mathclose{}\\right)\"\n )\n integration_utils.check_function(sum_with_limit, latex)\n\n\ndef test_sum_with_limit_2args() -> None:\n def sum_with_limit(a, n):\n return sum(i**2 for i in range(a, n))\n\n latex = (\n r\"\\mathrm{sum\\_with\\_limit}(a, n) = \\sum_{i = a}^{n - 1}\"\n r\" \\mathopen{}\\left({i^{2}}\\mathclose{}\\right)\"\n )\n integration_utils.check_function(sum_with_limit, latex)\n\n\ndef test_sum_with_reducible_limit() -> None:\n def sum_with_limit(n):\n return sum(i for i in range(n + 1))\n\n latex = (\n r\"\\mathrm{sum\\_with\\_limit}(n) = \\sum_{i = 0}^{n}\"\n r\" \\mathopen{}\\left({i}\\mathclose{}\\right)\"\n )\n integration_utils.check_function(sum_with_limit, latex)\n\n\ndef test_sum_with_irreducible_limit() -> None:\n def sum_with_limit(n):\n return sum(i for i in range(n * 3))\n\n latex = (\n r\"\\mathrm{sum\\_with\\_limit}(n) = \\sum_{i = 0}^{n \\cdot 3 - 1}\"\n r\" \\mathopen{}\\left({i}\\mathclose{}\\right)\"\n )\n integration_utils.check_function(sum_with_limit, latex)\n\n\ndef test_prod_with_limit_1arg() -> None:\n def prod_with_limit(n):\n return math.prod(i**2 for i in range(n))\n\n latex = (\n r\"\\mathrm{prod\\_with\\_limit}(n) =\"\n r\" \\prod_{i = 0}^{n - 1} \\mathopen{}\\left({i^{2}}\\mathclose{}\\right)\"\n )\n integration_utils.check_function(prod_with_limit, latex)\n\n\ndef test_prod_with_limit_2args() -> None:\n def prod_with_limit(a, n):\n return math.prod(i**2 for i in range(a, n))\n\n latex = (\n r\"\\mathrm{prod\\_with\\_limit}(a, n) =\"\n r\" \\prod_{i = a}^{n - 1} \\mathopen{}\\left({i^{2}}\\mathclose{}\\right)\"\n )\n integration_utils.check_function(prod_with_limit, latex)\n\n\ndef test_prod_with_reducible_limits() -> None:\n def prod_with_limit(n):\n return math.prod(i for i in range(n - 1))\n\n latex = (\n r\"\\mathrm{prod\\_with\\_limit}(n) =\"\n r\" \\prod_{i = 0}^{n - 2} \\mathopen{}\\left({i}\\mathclose{}\\right)\"\n )\n integration_utils.check_function(prod_with_limit, latex)\n\n\ndef test_prod_with_irreducible_limit() -> None:\n def prod_with_limit(n):\n return math.prod(i for i in range(n * 3))\n\n latex = (\n r\"\\mathrm{prod\\_with\\_limit}(n) = \"\n r\"\\prod_{i = 0}^{n \\cdot 3 - 1} \\mathopen{}\\left({i}\\mathclose{}\\right)\"\n )\n integration_utils.check_function(prod_with_limit, latex)\n\n\ndef test_nested_function() -> None:\n def nested(x):\n return 3 * x\n\n integration_utils.check_function(nested, r\"\\mathrm{nested}(x) = 3 x\")\n\n\ndef test_double_nested_function() -> None:\n def nested(x):\n def inner(y):\n return x * y\n\n return inner\n\n integration_utils.check_function(nested(3), r\"\\mathrm{inner}(y) = x y\")\n\n\ndef test_reduce_assignments() -> None:\n def f(x):\n a = x + x\n return 3 * a\n\n integration_utils.check_function(\n f,\n r\"\\begin{array}{l} a = x + x \\\\ f(x) = 3 a \\end{array}\",\n )\n integration_utils.check_function(\n f,\n r\"f(x) = 3 \\mathopen{}\\left( x + x \\mathclose{}\\right)\",\n reduce_assignments=True,\n )\n\n\ndef test_reduce_assignments_double() -> None:\n def f(x):\n a = x**2\n b = a + a\n return 3 * b\n\n latex_without_option = (\n r\"\\begin{array}{l}\"\n r\" a = x^{2} \\\\\"\n r\" b = a + a \\\\\"\n r\" f(x) = 3 b\"\n r\" \\end{array}\"\n )\n\n integration_utils.check_function(f, latex_without_option)\n integration_utils.check_function(f, latex_without_option, reduce_assignments=False)\n integration_utils.check_function(\n f,\n r\"f(x) = 3 \\mathopen{}\\left( x^{2} + x^{2} \\mathclose{}\\right)\",\n reduce_assignments=True,\n )\n\n\ndef test_reduce_assignments_with_if() -> None:\n def sigmoid(x):\n p = 1 / (1 + math.exp(-x))\n n = math.exp(x) / (math.exp(x) + 1)\n if x > 0:\n return p\n else:\n return n\n\n integration_utils.check_function(\n sigmoid,\n (\n r\"\\mathrm{sigmoid}(x) = \\left\\{ \\begin{array}{ll}\"\n r\" \\frac{1}{1 + \\exp \\mathopen{}\\left( -x \\mathclose{}\\right)}, &\"\n r\" \\mathrm{if} \\ x > 0 \\\\\"\n r\" \\frac{\\exp x}{\\exp x + 1}, &\"\n r\" \\mathrm{otherwise}\"\n r\" \\end{array} \\right.\"\n ),\n reduce_assignments=True,\n )\n\n\ndef test_sub_bracket() -> None:\n def solve(a, b):\n return ((a + b) - b) / (a - b) - (a + b) - (a - b) - (a * b)\n\n latex = (\n r\"\\mathrm{solve}(a, b) =\"\n r\" \\frac{a + b - b}{a - b} - \\mathopen{}\\left(\"\n r\" a + b \\mathclose{}\\right) - \\mathopen{}\\left(\"\n r\" a - b \\mathclose{}\\right) - a b\"\n )\n integration_utils.check_function(solve, latex)\n\n\ndef test_docstring_allowed() -> None:\n def solve(x):\n \"\"\"The identity function.\"\"\"\n return x\n\n latex = r\"\\mathrm{solve}(x) = x\"\n integration_utils.check_function(solve, latex)\n\n\ndef test_multiple_constants_allowed() -> None:\n def solve(x):\n \"\"\"The identity function.\"\"\"\n 123\n True\n return x\n\n latex = r\"\\mathrm{solve}(x) = x\"\n integration_utils.check_function(solve, latex)\n"
},
{
"path": "src/latexify/__init__.py",
"content": "\"\"\"Latexify root package.\"\"\"\n\ntry:\n from latexify import _version\n\n __version__ = _version.__version__\nexcept Exception:\n __version__ = \"\"\n\nfrom latexify import frontend, generate_latex\n\nStyle = generate_latex.Style\n\nget_latex = generate_latex.get_latex\n\nalgorithmic = frontend.algorithmic\nexpression = frontend.expression\nfunction = frontend.function\n"
},
{
"path": "src/latexify/_version.py",
"content": "\"\"\"Version specifier.\n\nDON'T TOUCH THIS FILE.\nThis file is replaced during the release process.\n\"\"\"\n\n__version__ = \"0.0.0a0\"\n"
},
{
"path": "src/latexify/analyzers.py",
"content": "\"\"\"Analyzer functions for specific subtrees.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\nimport dataclasses\nimport sys\n\nfrom latexify import ast_utils, exceptions\n\n\n@dataclasses.dataclass(frozen=True, eq=False)\nclass RangeInfo:\n \"\"\"Information of the range function.\"\"\"\n\n # Argument subtrees. These arguments could be shallow copies of the original\n # subtree.\n start: ast.expr\n stop: ast.expr\n step: ast.expr\n\n # Integer representation of each argument, when it is possible.\n start_int: int | None\n stop_int: int | None\n step_int: int | None\n\n\ndef analyze_range(node: ast.Call) -> RangeInfo:\n \"\"\"Obtains RangeInfo from a Call subtree.\n\n Args:\n node: Subtree to be analyzed.\n\n Returns:\n RangeInfo extracted from `node`.\n\n Raises:\n LatexifySyntaxError: Analysis failed.\n \"\"\"\n if not (\n isinstance(node.func, ast.Name)\n and node.func.id == \"range\"\n and 1 <= len(node.args) <= 3\n ):\n raise exceptions.LatexifySyntaxError(\"Unsupported AST for analyze_range.\")\n\n num_args = len(node.args)\n\n if num_args == 1:\n start = ast_utils.make_constant(0)\n stop = node.args[0]\n step = ast_utils.make_constant(1)\n else:\n start = node.args[0]\n stop = node.args[1]\n step = node.args[2] if num_args == 3 else ast_utils.make_constant(1)\n\n return RangeInfo(\n start=start,\n stop=stop,\n step=step,\n start_int=ast_utils.extract_int_or_none(start),\n stop_int=ast_utils.extract_int_or_none(stop),\n step_int=ast_utils.extract_int_or_none(step),\n )\n\n\ndef reduce_stop_parameter(node: ast.expr) -> ast.expr:\n \"\"\"Adjusts the stop expression of the range.\n\n This function tries to convert the syntax as follows:\n * n + 1 --> n\n * n + 2 --> n + 1\n * n - 1 --> n - 2\n\n Args:\n node: The target expression.\n\n Returns:\n Converted expression.\n \"\"\"\n if not (isinstance(node, ast.BinOp) and isinstance(node.op, (ast.Add, ast.Sub))):\n return ast.BinOp(left=node, op=ast.Sub(), right=ast_utils.make_constant(1))\n\n # Treatment for Python 3.7.\n rhs = (\n ast.Constant(value=node.right.n)\n if sys.version_info.minor < 8 and isinstance(node.right, ast.Num)\n else node.right\n )\n\n if not isinstance(rhs, ast.Constant):\n return ast.BinOp(left=node, op=ast.Sub(), right=ast_utils.make_constant(1))\n\n shift = 1 if isinstance(node.op, ast.Add) else -1\n\n return (\n node.left\n if rhs.value == shift\n else ast.BinOp(\n left=node.left,\n op=node.op,\n right=ast_utils.make_constant(value=rhs.value - shift),\n )\n )\n"
},
{
"path": "src/latexify/analyzers_test.py",
"content": "\"\"\"Tests for latexify.analyzers.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\n\nimport pytest\n\nfrom latexify import analyzers, ast_utils, exceptions, test_utils\n\n\n@pytest.mark.parametrize(\n \"code,start,stop,step,start_int,stop_int,step_int\",\n [\n (\n \"range(x)\",\n ast.Constant(value=0),\n ast.Name(id=\"x\", ctx=ast.Load()),\n ast.Constant(value=1),\n 0,\n None,\n 1,\n ),\n (\n \"range(123)\",\n ast.Constant(value=0),\n ast.Constant(value=123),\n ast.Constant(value=1),\n 0,\n 123,\n 1,\n ),\n (\n \"range(x, y)\",\n ast.Name(id=\"x\", ctx=ast.Load()),\n ast.Name(id=\"y\", ctx=ast.Load()),\n ast.Constant(value=1),\n None,\n None,\n 1,\n ),\n (\n \"range(123, y)\",\n ast.Constant(value=123),\n ast.Name(id=\"y\", ctx=ast.Load()),\n ast.Constant(value=1),\n 123,\n None,\n 1,\n ),\n (\n \"range(x, 123)\",\n ast.Name(id=\"x\", ctx=ast.Load()),\n ast.Constant(value=123),\n ast.Constant(value=1),\n None,\n 123,\n 1,\n ),\n (\n \"range(x, y, z)\",\n ast.Name(id=\"x\", ctx=ast.Load()),\n ast.Name(id=\"y\", ctx=ast.Load()),\n ast.Name(id=\"z\", ctx=ast.Load()),\n None,\n None,\n None,\n ),\n (\n \"range(123, y, z)\",\n ast.Constant(value=123),\n ast.Name(id=\"y\", ctx=ast.Load()),\n ast.Name(id=\"z\", ctx=ast.Load()),\n 123,\n None,\n None,\n ),\n (\n \"range(x, 123, z)\",\n ast.Name(id=\"x\", ctx=ast.Load()),\n ast.Constant(value=123),\n ast.Name(id=\"z\", ctx=ast.Load()),\n None,\n 123,\n None,\n ),\n (\n \"range(x, y, 123)\",\n ast.Name(id=\"x\", ctx=ast.Load()),\n ast.Name(id=\"y\", ctx=ast.Load()),\n ast.Constant(value=123),\n None,\n None,\n 123,\n ),\n ],\n)\ndef test_analyze_range(\n code: str,\n start: ast.expr,\n stop: ast.expr,\n step: ast.expr,\n start_int: int | None,\n stop_int: int | None,\n step_int: int | None,\n) -> None:\n node = ast_utils.parse_expr(code)\n assert isinstance(node, ast.Call)\n\n info = analyzers.analyze_range(node)\n\n test_utils.assert_ast_equal(observed=info.start, expected=start)\n test_utils.assert_ast_equal(observed=info.stop, expected=stop)\n if step is not None:\n test_utils.assert_ast_equal(observed=info.step, expected=step)\n else:\n assert info.step is None\n\n def check_int(observed: int | None, expected: int | None) -> None:\n if expected is not None:\n assert observed == expected\n else:\n assert observed is None\n\n check_int(observed=info.start_int, expected=start_int)\n check_int(observed=info.stop_int, expected=stop_int)\n check_int(observed=info.step_int, expected=step_int)\n\n\n@pytest.mark.parametrize(\n \"code\",\n [\n # Not a direct call\n \"__builtins__.range(x)\",\n 'getattr(__builtins__, \"range\")(x)',\n # Unsupported functions\n \"f(x)\",\n \"iter(range(x))\",\n # Range with invalid arguments\n \"range()\",\n \"range(x, y, z, w)\",\n ],\n)\ndef test_analyze_range_invalid(code: str) -> None:\n node = ast_utils.parse_expr(code)\n assert isinstance(node, ast.Call)\n\n with pytest.raises(\n exceptions.LatexifySyntaxError, match=r\"^Unsupported AST for analyze_range\\.$\"\n ):\n analyzers.analyze_range(node)\n\n\n@pytest.mark.parametrize(\n \"before,after\",\n [\n (\"n + 1\", \"n\"),\n (\"n + 2\", \"n + 1\"),\n (\"n - (-1)\", \"n - (-1) - 1\"),\n (\"n - 1\", \"n - 2\"),\n (\"1 * 2\", \"1 * 2 - 1\"),\n ],\n)\ndef test_reduce_stop_parameter(before: str, after: str) -> None:\n test_utils.assert_ast_equal(\n analyzers.reduce_stop_parameter(ast_utils.parse_expr(before)),\n ast_utils.parse_expr(after),\n )\n"
},
{
"path": "src/latexify/ast_utils.py",
"content": "\"\"\"Utilities to generate AST nodes.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\nimport sys\nfrom typing import Any\n\n\ndef parse_expr(code: str) -> ast.expr:\n \"\"\"Parses given Python expression.\n\n Args:\n code: Python expression to parse.\n\n Returns:\n ast.expr corresponding to `code`.\n \"\"\"\n return ast.parse(code, mode=\"eval\").body\n\n\ndef make_name(id: str) -> ast.Name:\n \"\"\"Generates a new Name node.\n\n Args:\n id: Name of the node.\n\n Returns:\n Generated ast.Name.\n \"\"\"\n return ast.Name(id=id, ctx=ast.Load())\n\n\ndef make_attribute(value: ast.expr, attr: str):\n \"\"\"Generates a new Attribute node.\n\n Args:\n value: Parent value.\n attr: Attribute name.\n\n Returns:\n Generated ast.Attribute.\n \"\"\"\n return ast.Attribute(value=value, attr=attr, ctx=ast.Load())\n\n\ndef make_constant(value: Any) -> ast.expr:\n \"\"\"Generates a new Constant node.\n\n Args:\n value: Value of the node.\n\n Returns:\n Generated ast.Constant or its equivalent.\n\n Raises:\n ValueError: Unsupported value type.\n \"\"\"\n if (\n value is None\n or value is ...\n or isinstance(value, (bool, int, float, complex, str, bytes))\n ):\n return ast.Constant(value=value)\n\n raise ValueError(f\"Unsupported type to generate Constant: {type(value).__name__}\")\n\n\ndef is_constant(node: ast.AST) -> bool:\n \"\"\"Checks if the node is a constant.\n\n Args:\n node: The node to examine.\n\n Returns:\n True if the node is a constant, False otherwise.\n \"\"\"\n return isinstance(node, ast.Constant)\n\n\ndef is_str(node: ast.AST) -> bool:\n \"\"\"Checks if the node is a str constant.\n\n Args:\n node: The node to examine.\n\n Returns:\n True if the node is a str constant, False otherwise.\n \"\"\"\n if sys.version_info.minor < 8 and isinstance(node, ast.Str):\n return True\n\n return isinstance(node, ast.Constant) and isinstance(node.value, str)\n\n\ndef extract_int_or_none(node: ast.expr) -> int | None:\n \"\"\"Extracts int constant from the given Constant node.\n\n Args:\n node: ast.Constant or its equivalent representing an int value.\n\n Returns:\n Extracted int value, or None if extraction failed.\n \"\"\"\n if (\n isinstance(node, ast.Constant)\n and isinstance(node.value, int)\n and not isinstance(node.value, bool)\n ):\n return node.value\n\n return None\n\n\ndef extract_int(node: ast.expr) -> int:\n \"\"\"Extracts int constant from the given Constant node.\n\n Args:\n node: ast.Constant or its equivalent representing an int value.\n\n Returns:\n Extracted int value.\n\n Raises:\n ValueError: Not a subtree containing an int value.\n \"\"\"\n value = extract_int_or_none(node)\n\n if value is None:\n raise ValueError(f\"Unsupported node to extract int: {type(node).__name__}\")\n\n return value\n\n\ndef extract_function_name_or_none(node: ast.Call) -> str | None:\n \"\"\"Extracts function name from the given Call node.\n\n Args:\n node: ast.Call.\n\n Returns:\n Extracted function name, or None if not found.\n \"\"\"\n if isinstance(node.func, ast.Name):\n return node.func.id\n if isinstance(node.func, ast.Attribute):\n return node.func.attr\n\n return None\n\n\ndef create_function_def(\n name,\n args,\n body,\n decorator_list,\n returns=None,\n type_comment=None,\n type_params=None,\n lineno=None,\n col_offset=None,\n end_lineno=None,\n end_col_offset=None,\n) -> ast.FunctionDef:\n \"\"\"Creates a FunctionDef node.\n\n This function generates an `ast.FunctionDef` node, optionally removing\n the `type_params` keyword argument for Python versions below 3.12.\n\n Args:\n name: Name of the function.\n args: Arguments of the function.\n body: Body of the function.\n decorator_list: List of decorators.\n returns: Return type of the function.\n type_comment: Type comment of the function.\n type_params: Type parameters of the function.\n lineno: Line number of the function definition.\n col_offset: Column offset of the function definition.\n end_lineno: End line number of the function definition.\n end_col_offset: End column offset of the function definition.\n\n Returns:\n ast.FunctionDef: The generated FunctionDef node.\n \"\"\"\n if sys.version_info.minor < 12:\n return ast.FunctionDef(\n name=name,\n args=args,\n body=body,\n decorator_list=decorator_list,\n returns=returns,\n type_comment=type_comment,\n lineno=lineno,\n col_offset=col_offset,\n end_lineno=end_lineno,\n end_col_offset=end_col_offset,\n ) # type: ignore\n return ast.FunctionDef(\n name=name,\n args=args,\n body=body,\n decorator_list=decorator_list,\n returns=returns,\n type_comment=type_comment,\n type_params=type_params,\n lineno=lineno,\n col_offset=col_offset,\n end_lineno=end_lineno,\n end_col_offset=end_col_offset,\n ) # type: ignore\n"
},
{
"path": "src/latexify/ast_utils_test.py",
"content": "\"\"\"Tests for latexify.ast_utils.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\nimport sys\nfrom typing import Any\n\nimport pytest\n\nfrom latexify import ast_utils, test_utils\n\n\ndef test_parse_expr() -> None:\n test_utils.assert_ast_equal(\n ast_utils.parse_expr(\"a + b\"),\n ast.BinOp(\n left=ast_utils.make_name(\"a\"),\n op=ast.Add(),\n right=ast_utils.make_name(\"b\"),\n ),\n )\n\n\ndef test_make_name() -> None:\n test_utils.assert_ast_equal(\n ast_utils.make_name(\"foo\"), ast.Name(id=\"foo\", ctx=ast.Load())\n )\n\n\ndef test_make_attribute() -> None:\n test_utils.assert_ast_equal(\n ast_utils.make_attribute(ast_utils.make_name(\"foo\"), \"bar\"),\n ast.Attribute(ast.Name(id=\"foo\", ctx=ast.Load()), attr=\"bar\", ctx=ast.Load()),\n )\n\n\n@pytest.mark.parametrize(\n \"value,expected\",\n [\n (None, ast.Constant(value=None)),\n (False, ast.Constant(value=False)),\n (True, ast.Constant(value=True)),\n (..., ast.Constant(value=...)),\n (123, ast.Constant(value=123)),\n (4.5, ast.Constant(value=4.5)),\n (6 + 7j, ast.Constant(value=6 + 7j)),\n (\"foo\", ast.Constant(value=\"foo\")),\n (b\"bar\", ast.Constant(value=b\"bar\")),\n ],\n)\ndef test_make_constant(value: Any, expected: ast.Constant) -> None:\n test_utils.assert_ast_equal(\n observed=ast_utils.make_constant(value),\n expected=expected,\n )\n\n\ndef test_make_constant_invalid() -> None:\n with pytest.raises(ValueError, match=r\"^Unsupported type to generate\"):\n ast_utils.make_constant(object())\n\n\n@pytest.mark.parametrize(\n \"value,expected\",\n [\n (ast.Constant(value=\"foo\"), True),\n (ast.Expr(value=ast.Constant(value=123)), False),\n (ast.Global(names=[\"bar\"]), False),\n ],\n)\ndef test_is_constant(value: ast.AST, expected: bool) -> None:\n assert ast_utils.is_constant(value) is expected\n\n\n@pytest.mark.parametrize(\n \"value,expected\",\n [\n (ast.Constant(value=123), False),\n (ast.Constant(value=\"foo\"), True),\n (ast.Expr(value=ast.Constant(value=\"foo\")), False),\n (ast.Global(names=[\"foo\"]), False),\n ],\n)\ndef test_is_str(value: ast.AST, expected: bool) -> None:\n assert ast_utils.is_str(value) is expected\n\n\ndef test_extract_int_or_none() -> None:\n assert ast_utils.extract_int_or_none(ast_utils.make_constant(-123)) == -123\n assert ast_utils.extract_int_or_none(ast_utils.make_constant(0)) == 0\n assert ast_utils.extract_int_or_none(ast_utils.make_constant(123)) == 123\n\n\ndef test_extract_int_or_none_invalid() -> None:\n # Not a Constant node with int.\n assert ast_utils.extract_int_or_none(ast_utils.make_constant(None)) is None\n assert ast_utils.extract_int_or_none(ast_utils.make_constant(True)) is None\n assert ast_utils.extract_int_or_none(ast_utils.make_constant(...)) is None\n assert ast_utils.extract_int_or_none(ast_utils.make_constant(123.0)) is None\n assert ast_utils.extract_int_or_none(ast_utils.make_constant(4 + 5j)) is None\n assert ast_utils.extract_int_or_none(ast_utils.make_constant(\"123\")) is None\n assert ast_utils.extract_int_or_none(ast_utils.make_constant(b\"123\")) is None\n\n\ndef test_extract_int() -> None:\n assert ast_utils.extract_int(ast_utils.make_constant(-123)) == -123\n assert ast_utils.extract_int(ast_utils.make_constant(0)) == 0\n assert ast_utils.extract_int(ast_utils.make_constant(123)) == 123\n\n\ndef test_extract_int_invalid() -> None:\n # Not a Constant node with int.\n with pytest.raises(ValueError, match=r\"^Unsupported node to extract int\"):\n ast_utils.extract_int(ast_utils.make_constant(None))\n with pytest.raises(ValueError, match=r\"^Unsupported node to extract int\"):\n ast_utils.extract_int(ast_utils.make_constant(True))\n with pytest.raises(ValueError, match=r\"^Unsupported node to extract int\"):\n ast_utils.extract_int(ast_utils.make_constant(...))\n with pytest.raises(ValueError, match=r\"^Unsupported node to extract int\"):\n ast_utils.extract_int(ast_utils.make_constant(123.0))\n with pytest.raises(ValueError, match=r\"^Unsupported node to extract int\"):\n ast_utils.extract_int(ast_utils.make_constant(4 + 5j))\n with pytest.raises(ValueError, match=r\"^Unsupported node to extract int\"):\n ast_utils.extract_int(ast_utils.make_constant(\"123\"))\n with pytest.raises(ValueError, match=r\"^Unsupported node to extract int\"):\n ast_utils.extract_int(ast_utils.make_constant(b\"123\"))\n\n\n@pytest.mark.parametrize(\n \"value,expected\",\n [\n (\n ast.Call(\n func=ast.Name(id=\"hypot\", ctx=ast.Load()),\n args=[],\n keywords=[],\n ),\n \"hypot\",\n ),\n (\n ast.Call(\n func=ast.Attribute(\n value=ast.Name(id=\"math\", ctx=ast.Load()),\n attr=\"hypot\",\n ctx=ast.Load(),\n ),\n args=[],\n keywords=[],\n ),\n \"hypot\",\n ),\n (\n ast.Call(\n func=ast.Call(\n func=ast.Name(id=\"foo\", ctx=ast.Load()), args=[], keywords=[]\n ),\n args=[],\n keywords=[],\n ),\n None,\n ),\n ],\n)\ndef test_extract_function_name_or_none(value: ast.Call, expected: str | None) -> None:\n assert ast_utils.extract_function_name_or_none(value) == expected\n\n\ndef test_create_function_def() -> None:\n expected_args = ast.arguments(\n posonlyargs=[],\n args=[ast.arg(arg=\"x\")],\n vararg=None,\n kwonlyargs=[],\n kw_defaults=[],\n kwarg=None,\n defaults=[],\n )\n\n kwargs = {\n \"name\": \"test_func\",\n \"args\": expected_args,\n \"body\": [ast.Return(value=ast.Name(id=\"x\", ctx=ast.Load()))],\n \"decorator_list\": [],\n \"returns\": None,\n \"type_comment\": None,\n \"lineno\": 1,\n \"col_offset\": 0,\n \"end_lineno\": 2,\n \"end_col_offset\": 0,\n }\n if sys.version_info.minor >= 12:\n kwargs[\"type_params\"] = []\n\n func_def = ast_utils.create_function_def(**kwargs)\n assert isinstance(func_def, ast.FunctionDef)\n assert func_def.name == \"test_func\"\n\n assert func_def.args.posonlyargs == expected_args.posonlyargs\n assert func_def.args.args == expected_args.args\n assert func_def.args.kwonlyargs == expected_args.kwonlyargs\n assert func_def.args.kw_defaults == expected_args.kw_defaults\n assert func_def.args.defaults == expected_args.defaults\n"
},
{
"path": "src/latexify/codegen/__init__.py",
"content": "\"\"\"Package latexify.codegen.\"\"\"\n\nfrom latexify.codegen import algorithmic_codegen, expression_codegen, function_codegen\n\nAlgorithmicCodegen = algorithmic_codegen.AlgorithmicCodegen\nExpressionCodegen = expression_codegen.ExpressionCodegen\nFunctionCodegen = function_codegen.FunctionCodegen\nIPythonAlgorithmicCodegen = algorithmic_codegen.IPythonAlgorithmicCodegen\n"
},
{
"path": "src/latexify/codegen/algorithmic_codegen.py",
"content": "\"\"\"Codegen for single algorithms.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\nimport contextlib\nfrom collections.abc import Generator\n\nfrom latexify import exceptions\nfrom latexify.codegen import expression_codegen, identifier_converter\n\n\nclass AlgorithmicCodegen(ast.NodeVisitor):\n \"\"\"Codegen for single algorithms.\n\n This codegen works for Module with single FunctionDef node to generate a single\n LaTeX expression of the given algorithm.\n \"\"\"\n\n _SPACES_PER_INDENT = 4\n\n _identifier_converter: identifier_converter.IdentifierConverter\n _indent_level: int\n\n def __init__(\n self,\n *,\n use_math_symbols: bool = False,\n use_set_symbols: bool = False,\n escape_underscores: bool = True,\n ) -> None:\n \"\"\"Initializer.\n\n Args:\n use_math_symbols: Whether to convert identifiers with a math symbol surface\n (e.g., \"alpha\") to the LaTeX symbol (e.g., \"\\\\alpha\").\n use_set_symbols: Whether to use set symbols or not.\n \"\"\"\n self._expression_codegen = expression_codegen.ExpressionCodegen(\n use_math_symbols=use_math_symbols,\n use_set_symbols=use_set_symbols,\n escape_underscores=escape_underscores,\n )\n self._identifier_converter = identifier_converter.IdentifierConverter(\n use_math_symbols=use_math_symbols,\n use_mathrm=False,\n escape_underscores=escape_underscores,\n )\n self._indent_level = 0\n\n def generic_visit(self, node: ast.AST) -> str:\n raise exceptions.LatexifyNotSupportedError(\n f\"Unsupported AST: {type(node).__name__}\"\n )\n\n def visit_Assign(self, node: ast.Assign) -> str:\n \"\"\"Visit an Assign node.\"\"\"\n operands: list[str] = [\n self._expression_codegen.visit(target) for target in node.targets\n ]\n operands.append(self._expression_codegen.visit(node.value))\n operands_latex = r\" \\gets \".join(operands)\n return self._add_indent(rf\"\\State ${operands_latex}$\")\n\n def visit_Expr(self, node: ast.Expr) -> str:\n \"\"\"Visit an Expr node.\"\"\"\n return self._add_indent(\n rf\"\\State ${self._expression_codegen.visit(node.value)}$\"\n )\n\n def visit_For(self, node: ast.For) -> str:\n \"\"\"Visit a For node.\"\"\"\n if len(node.orelse) != 0:\n raise exceptions.LatexifyNotSupportedError(\n \"For statement with the else clause is not supported\"\n )\n\n target_latex = self._expression_codegen.visit(node.target)\n iter_latex = self._expression_codegen.visit(node.iter)\n with self._increment_level():\n body_latex = \"\\n\".join(self.visit(stmt) for stmt in node.body)\n\n return (\n self._add_indent(f\"\\\\For{{${target_latex} \\\\in {iter_latex}$}}\\n\")\n + f\"{body_latex}\\n\"\n + self._add_indent(\"\\\\EndFor\")\n )\n\n # TODO(ZibingZhang): support nested functions\n def visit_FunctionDef(self, node: ast.FunctionDef) -> str:\n \"\"\"Visit a FunctionDef node.\"\"\"\n name_latex = self._identifier_converter.convert(node.name)[0]\n\n # Arguments\n arg_strs = [\n self._identifier_converter.convert(arg.arg)[0] for arg in node.args.args\n ]\n\n latex = self._add_indent(\"\\\\begin{algorithmic}\\n\")\n with self._increment_level():\n latex += self._add_indent(\n f\"\\\\Function{{{name_latex}}}{{${', '.join(arg_strs)}$}}\\n\"\n )\n\n with self._increment_level():\n # Body\n body_strs: list[str] = [self.visit(stmt) for stmt in node.body]\n body_latex = \"\\n\".join(body_strs)\n\n latex += f\"{body_latex}\\n\"\n latex += self._add_indent(\"\\\\EndFunction\\n\")\n return latex + self._add_indent(r\"\\end{algorithmic}\")\n\n # TODO(ZibingZhang): support \\ELSIF\n def visit_If(self, node: ast.If) -> str:\n \"\"\"Visit an If node.\"\"\"\n cond_latex = self._expression_codegen.visit(node.test)\n with self._increment_level():\n body_latex = \"\\n\".join(self.visit(stmt) for stmt in node.body)\n\n latex = self._add_indent(f\"\\\\If{{${cond_latex}$}}\\n\" + body_latex)\n\n if node.orelse:\n latex += \"\\n\" + self._add_indent(\"\\\\Else\\n\")\n with self._increment_level():\n latex += \"\\n\".join(self.visit(stmt) for stmt in node.orelse)\n\n return f\"{latex}\\n\" + self._add_indent(r\"\\EndIf\")\n\n def visit_Module(self, node: ast.Module) -> str:\n \"\"\"Visit a Module node.\"\"\"\n return self.visit(node.body[0])\n\n def visit_Return(self, node: ast.Return) -> str:\n \"\"\"Visit a Return node.\"\"\"\n return (\n self._add_indent(\n rf\"\\State \\Return ${self._expression_codegen.visit(node.value)}$\"\n )\n if node.value is not None\n else self._add_indent(r\"\\State \\Return\")\n )\n\n def visit_While(self, node: ast.While) -> str:\n \"\"\"Visit a While node.\"\"\"\n if node.orelse:\n raise exceptions.LatexifyNotSupportedError(\n \"While statement with the else clause is not supported\"\n )\n\n cond_latex = self._expression_codegen.visit(node.test)\n with self._increment_level():\n body_latex = \"\\n\".join(self.visit(stmt) for stmt in node.body)\n return (\n self._add_indent(f\"\\\\While{{${cond_latex}$}}\\n\")\n + f\"{body_latex}\\n\"\n + self._add_indent(r\"\\EndWhile\")\n )\n\n def visit_Pass(self, node: ast.Pass) -> str:\n \"\"\"Visit a Pass node.\"\"\"\n return self._add_indent(r\"\\State $\\mathbf{pass}$\")\n\n def visit_Break(self, node: ast.Break) -> str:\n \"\"\"Visit a Break node.\"\"\"\n return self._add_indent(r\"\\State $\\mathbf{break}$\")\n\n def visit_Continue(self, node: ast.Continue) -> str:\n \"\"\"Visit a Continue node.\"\"\"\n return self._add_indent(r\"\\State $\\mathbf{continue}$\")\n\n @contextlib.contextmanager\n def _increment_level(self) -> Generator[None, None, None]:\n \"\"\"Context manager controlling indent level.\"\"\"\n self._indent_level += 1\n yield\n self._indent_level -= 1\n\n def _add_indent(self, line: str) -> str:\n \"\"\"Adds an indent before the line.\n\n Args:\n line: The line to add an indent to.\n \"\"\"\n return self._indent_level * self._SPACES_PER_INDENT * \" \" + line\n\n\nclass IPythonAlgorithmicCodegen(ast.NodeVisitor):\n \"\"\"Codegen for single algorithms targeting IPython.\n\n This codegen works for Module with single FunctionDef node to generate a single\n LaTeX expression of the given algorithm.\n \"\"\"\n\n _EM_PER_INDENT = 1\n _LINE_BREAK = r\" \\\\ \"\n\n _identifier_converter: identifier_converter.IdentifierConverter\n _indent_level: int\n\n def __init__(\n self,\n *,\n use_math_symbols: bool = False,\n use_set_symbols: bool = False,\n escape_underscores: bool = True,\n ) -> None:\n \"\"\"Initializer.\n\n Args:\n use_math_symbols: Whether to convert identifiers with a math symbol surface\n (e.g., \"alpha\") to the LaTeX symbol (e.g., \"\\\\alpha\").\n use_set_symbols: Whether to use set symbols or not.\n \"\"\"\n self._expression_codegen = expression_codegen.ExpressionCodegen(\n use_math_symbols=use_math_symbols,\n use_set_symbols=use_set_symbols,\n escape_underscores=escape_underscores,\n )\n self._identifier_converter = identifier_converter.IdentifierConverter(\n use_math_symbols=use_math_symbols, escape_underscores=escape_underscores\n )\n self._indent_level = 0\n\n def generic_visit(self, node: ast.AST) -> str:\n raise exceptions.LatexifyNotSupportedError(\n f\"Unsupported AST: {type(node).__name__}\"\n )\n\n def visit_Assign(self, node: ast.Assign) -> str:\n \"\"\"Visit an Assign node.\"\"\"\n operands: list[str] = [\n self._expression_codegen.visit(target) for target in node.targets\n ]\n operands.append(self._expression_codegen.visit(node.value))\n operands_latex = r\" \\gets \".join(operands)\n return self._add_indent(operands_latex)\n\n def visit_Expr(self, node: ast.Expr) -> str:\n \"\"\"Visit an Expr node.\"\"\"\n return self._add_indent(self._expression_codegen.visit(node.value))\n\n def visit_For(self, node: ast.For) -> str:\n \"\"\"Visit a For node.\"\"\"\n if len(node.orelse) != 0:\n raise exceptions.LatexifyNotSupportedError(\n \"For statement with the else clause is not supported\"\n )\n\n target_latex = self._expression_codegen.visit(node.target)\n iter_latex = self._expression_codegen.visit(node.iter)\n with self._increment_level():\n body_latex = self._LINE_BREAK.join(self.visit(stmt) for stmt in node.body)\n\n return (\n self._add_indent(r\"\\mathbf{for}\")\n + rf\" \\ {target_latex} \\in {iter_latex} \\ \\mathbf{{do}}{self._LINE_BREAK}\"\n + f\"{body_latex}{self._LINE_BREAK}\"\n + self._add_indent(r\"\\mathbf{end \\ for}\")\n )\n\n # TODO(ZibingZhang): support nested functions\n def visit_FunctionDef(self, node: ast.FunctionDef) -> str:\n \"\"\"Visit a FunctionDef node.\"\"\"\n name_latex = self._identifier_converter.convert(node.name)[0]\n\n # Arguments\n args_latex = [\n self._identifier_converter.convert(arg.arg)[0] for arg in node.args.args\n ]\n # Body\n with self._increment_level():\n body_stmts_latex: list[str] = [self.visit(stmt) for stmt in node.body]\n body_latex = self._LINE_BREAK.join(body_stmts_latex)\n\n return (\n r\"\\begin{array}{l} \"\n + self._add_indent(r\"\\mathbf{function}\")\n + rf\" \\ {name_latex}({', '.join(args_latex)})\"\n + f\"{self._LINE_BREAK}{body_latex}{self._LINE_BREAK}\"\n + self._add_indent(r\"\\mathbf{end \\ function}\")\n + r\" \\end{array}\"\n )\n\n # TODO(ZibingZhang): support \\ELSIF\n def visit_If(self, node: ast.If) -> str:\n \"\"\"Visit an If node.\"\"\"\n cond_latex = self._expression_codegen.visit(node.test)\n with self._increment_level():\n body_latex = self._LINE_BREAK.join(self.visit(stmt) for stmt in node.body)\n latex = self._add_indent(\n rf\"\\mathbf{{if}} \\ {cond_latex}{self._LINE_BREAK}{body_latex}\"\n )\n\n if node.orelse:\n latex += self._LINE_BREAK + self._add_indent(r\"\\mathbf{else} \\\\ \")\n with self._increment_level():\n latex += self._LINE_BREAK.join(self.visit(stmt) for stmt in node.orelse)\n\n return latex + self._LINE_BREAK + self._add_indent(r\"\\mathbf{end \\ if}\")\n\n def visit_Module(self, node: ast.Module) -> str:\n \"\"\"Visit a Module node.\"\"\"\n return self.visit(node.body[0])\n\n def visit_Return(self, node: ast.Return) -> str:\n \"\"\"Visit a Return node.\"\"\"\n return (\n self._add_indent(r\"\\mathbf{return} \\ \")\n + self._expression_codegen.visit(node.value)\n if node.value is not None\n else self._add_indent(r\"\\mathbf{return}\")\n )\n\n def visit_While(self, node: ast.While) -> str:\n \"\"\"Visit a While node.\"\"\"\n if node.orelse:\n raise exceptions.LatexifyNotSupportedError(\n \"While statement with the else clause is not supported\"\n )\n\n cond_latex = self._expression_codegen.visit(node.test)\n with self._increment_level():\n body_latex = self._LINE_BREAK.join(self.visit(stmt) for stmt in node.body)\n return (\n self._add_indent(r\"\\mathbf{while} \\ \")\n + f\"{cond_latex}{self._LINE_BREAK}{body_latex}{self._LINE_BREAK}\"\n + self._add_indent(r\"\\mathbf{end \\ while}\")\n )\n\n def visit_Pass(self, node: ast.Pass) -> str:\n \"\"\"Visit a Pass node.\"\"\"\n return self._add_indent(r\"\\mathbf{pass}\")\n\n def visit_Break(self, node: ast.Break) -> str:\n \"\"\"Visit a Break node.\"\"\"\n return self._add_indent(r\"\\mathbf{break}\")\n\n def visit_Continue(self, node: ast.Continue) -> str:\n \"\"\"Visit a Continue node.\"\"\"\n return self._add_indent(r\"\\mathbf{continue}\")\n\n @contextlib.contextmanager\n def _increment_level(self) -> Generator[None, None, None]:\n \"\"\"Context manager controlling indent level.\"\"\"\n self._indent_level += 1\n yield\n self._indent_level -= 1\n\n def _add_indent(self, line: str) -> str:\n \"\"\"Adds an indent before the line.\n\n Args:\n line: The line to add an indent to.\n \"\"\"\n return (\n rf\"\\hspace{{{self._indent_level * self._EM_PER_INDENT}em}} {line}\"\n if self._indent_level > 0\n else line\n )\n"
},
{
"path": "src/latexify/codegen/algorithmic_codegen_test.py",
"content": "\"\"\"Tests for latexify.codegen.algorithmic_codegen.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\nimport textwrap\n\nimport pytest\n\nfrom latexify import exceptions\nfrom latexify.codegen import algorithmic_codegen\n\n\ndef test_generic_visit() -> None:\n class UnknownNode(ast.AST):\n pass\n\n with pytest.raises(\n exceptions.LatexifyNotSupportedError,\n match=r\"^Unsupported AST: UnknownNode$\",\n ):\n algorithmic_codegen.AlgorithmicCodegen().visit(UnknownNode())\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\n \"x = 3\",\n r\"\\State $x \\gets 3$\",\n ),\n (\n \"a = b = 0\",\n r\"\\State $a \\gets b \\gets 0$\",\n ),\n ],\n)\ndef test_visit_assign(code: str, latex: str) -> None:\n node = ast.parse(textwrap.dedent(code)).body[0]\n assert isinstance(node, ast.Assign)\n assert algorithmic_codegen.AlgorithmicCodegen().visit(node) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\n \"for i in {1}: x = i\",\n r\"\"\"\n \\For{$i \\in \\mathopen{}\\left\\{ 1 \\mathclose{}\\right\\}$}\n \\State $x \\gets i$\n \\EndFor\n \"\"\",\n ),\n ],\n)\ndef test_visit_for(code: str, latex: str) -> None:\n node = ast.parse(textwrap.dedent(code)).body[0]\n assert isinstance(node, ast.For)\n assert (\n algorithmic_codegen.AlgorithmicCodegen().visit(node)\n == textwrap.dedent(latex).strip()\n )\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\n \"def f(x): return x\",\n r\"\"\"\n \\begin{algorithmic}\n \\Function{f}{$x$}\n \\State \\Return $x$\n \\EndFunction\n \\end{algorithmic}\n \"\"\",\n ),\n (\n \"def xyz(a, b, c): return 3\",\n r\"\"\"\n \\begin{algorithmic}\n \\Function{xyz}{$a, b, c$}\n \\State \\Return $3$\n \\EndFunction\n \\end{algorithmic}\n \"\"\",\n ),\n ],\n)\ndef test_visit_functiondef(code: str, latex: str) -> None:\n node = ast.parse(textwrap.dedent(code)).body[0]\n assert isinstance(node, ast.FunctionDef)\n assert (\n algorithmic_codegen.AlgorithmicCodegen().visit(node)\n == textwrap.dedent(latex).strip()\n )\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\n \"if x < y: return x\",\n r\"\"\"\n \\If{$x < y$}\n \\State \\Return $x$\n \\EndIf\n \"\"\",\n ),\n (\n \"if True: x\\nelse: y\",\n r\"\"\"\n \\If{$\\mathrm{True}$}\n \\State $x$\n \\Else\n \\State $y$\n \\EndIf\n \"\"\",\n ),\n ],\n)\ndef test_visit_if(code: str, latex: str) -> None:\n node = ast.parse(code).body[0]\n assert isinstance(node, ast.If)\n assert (\n algorithmic_codegen.AlgorithmicCodegen().visit(node)\n == textwrap.dedent(latex).strip()\n )\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\n \"return x + y\",\n r\"\\State \\Return $x + y$\",\n ),\n (\n \"return\",\n r\"\\State \\Return\",\n ),\n ],\n)\ndef test_visit_return(code: str, latex: str) -> None:\n node = ast.parse(code).body[0]\n assert isinstance(node, ast.Return)\n assert algorithmic_codegen.AlgorithmicCodegen().visit(node) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\n \"while x < y: x = x + 1\",\n r\"\"\"\n \\While{$x < y$}\n \\State $x \\gets x + 1$\n \\EndWhile\n \"\"\",\n )\n ],\n)\ndef test_visit_while(code: str, latex: str) -> None:\n node = ast.parse(code).body[0]\n assert isinstance(node, ast.While)\n assert (\n algorithmic_codegen.AlgorithmicCodegen().visit(node)\n == textwrap.dedent(latex).strip()\n )\n\n\ndef test_visit_while_with_else() -> None:\n node = ast.parse(\n textwrap.dedent(\n \"\"\"\n while True:\n x = x\n else:\n x = y\n \"\"\"\n )\n ).body[0]\n assert isinstance(node, ast.While)\n with pytest.raises(\n exceptions.LatexifyNotSupportedError,\n match=\"^While statement with the else clause is not supported$\",\n ):\n algorithmic_codegen.AlgorithmicCodegen().visit(node)\n\n\ndef test_visit_pass() -> None:\n node = ast.parse(\"pass\").body[0]\n assert isinstance(node, ast.Pass)\n assert (\n algorithmic_codegen.AlgorithmicCodegen().visit(node)\n == r\"\\State $\\mathbf{pass}$\"\n )\n\n\ndef test_visit_break() -> None:\n node = ast.parse(\"break\").body[0]\n assert isinstance(node, ast.Break)\n assert (\n algorithmic_codegen.AlgorithmicCodegen().visit(node)\n == r\"\\State $\\mathbf{break}$\"\n )\n\n\ndef test_visit_continue() -> None:\n node = ast.parse(\"continue\").body[0]\n assert isinstance(node, ast.Continue)\n assert (\n algorithmic_codegen.AlgorithmicCodegen().visit(node)\n == r\"\\State $\\mathbf{continue}$\"\n )\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"x = 3\", r\"x \\gets 3\"),\n (\"a = b = 0\", r\"a \\gets b \\gets 0\"),\n ],\n)\ndef test_visit_assign_ipython(code: str, latex: str) -> None:\n node = ast.parse(textwrap.dedent(code)).body[0]\n assert isinstance(node, ast.Assign)\n assert algorithmic_codegen.IPythonAlgorithmicCodegen().visit(node) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\n \"for i in {1}: x = i\",\n (\n r\"\\mathbf{for} \\ i \\in \\mathopen{}\\left\\{ 1 \\mathclose{}\\right\\}\"\n r\" \\ \\mathbf{do} \\\\\"\n r\" \\hspace{1em} x \\gets i \\\\\"\n r\" \\mathbf{end \\ for}\"\n ),\n ),\n ],\n)\ndef test_visit_for_ipython(code: str, latex: str) -> None:\n node = ast.parse(textwrap.dedent(code)).body[0]\n assert isinstance(node, ast.For)\n assert (\n algorithmic_codegen.IPythonAlgorithmicCodegen().visit(node)\n == textwrap.dedent(latex).strip()\n )\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\n \"def f(x): return x\",\n (\n r\"\\begin{array}{l}\"\n r\" \\mathbf{function}\"\n r\" \\ f(x) \\\\\"\n r\" \\hspace{1em} \\mathbf{return} \\ x \\\\\"\n r\" \\mathbf{end \\ function}\"\n r\" \\end{array}\"\n ),\n ),\n (\n \"def f(a, b, c): return 3\",\n (\n r\"\\begin{array}{l}\"\n r\" \\mathbf{function}\"\n r\" \\ f(a, b, c) \\\\\"\n r\" \\hspace{1em} \\mathbf{return} \\ 3 \\\\\"\n r\" \\mathbf{end \\ function}\"\n r\" \\end{array}\"\n ),\n ),\n ],\n)\ndef test_visit_functiondef_ipython(code: str, latex: str) -> None:\n node = ast.parse(textwrap.dedent(code)).body[0]\n assert isinstance(node, ast.FunctionDef)\n assert algorithmic_codegen.IPythonAlgorithmicCodegen().visit(node) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\n \"if x < y: return x\",\n (\n r\"\\mathbf{if} \\ x < y \\\\\"\n r\" \\hspace{1em} \\mathbf{return} \\ x \\\\\"\n r\" \\mathbf{end \\ if}\"\n ),\n ),\n (\n \"if True: x\\nelse: y\",\n (\n r\"\\mathbf{if} \\ \\mathrm{True} \\\\\"\n r\" \\hspace{1em} x \\\\\"\n r\" \\mathbf{else} \\\\\"\n r\" \\hspace{1em} y \\\\\"\n r\" \\mathbf{end \\ if}\"\n ),\n ),\n ],\n)\ndef test_visit_if_ipython(code: str, latex: str) -> None:\n node = ast.parse(textwrap.dedent(code)).body[0]\n assert isinstance(node, ast.If)\n assert algorithmic_codegen.IPythonAlgorithmicCodegen().visit(node) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\n \"return x + y\",\n r\"\\mathbf{return} \\ x + y\",\n ),\n (\n \"return\",\n r\"\\mathbf{return}\",\n ),\n ],\n)\ndef test_visit_return_ipython(code: str, latex: str) -> None:\n node = ast.parse(textwrap.dedent(code)).body[0]\n assert isinstance(node, ast.Return)\n assert algorithmic_codegen.IPythonAlgorithmicCodegen().visit(node) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\n \"while x < y: x = x + 1\",\n (\n r\"\\mathbf{while} \\ x < y \\\\\"\n r\" \\hspace{1em} x \\gets x + 1 \\\\\"\n r\" \\mathbf{end \\ while}\"\n ),\n )\n ],\n)\ndef test_visit_while_ipython(code: str, latex: str) -> None:\n node = ast.parse(textwrap.dedent(code)).body[0]\n assert isinstance(node, ast.While)\n assert algorithmic_codegen.IPythonAlgorithmicCodegen().visit(node) == latex\n\n\ndef test_visit_while_with_else_ipython() -> None:\n node = ast.parse(\n textwrap.dedent(\n \"\"\"\n while True:\n x = x\n else:\n x = y\n \"\"\"\n )\n ).body[0]\n assert isinstance(node, ast.While)\n with pytest.raises(\n exceptions.LatexifyNotSupportedError,\n match=\"^While statement with the else clause is not supported$\",\n ):\n algorithmic_codegen.IPythonAlgorithmicCodegen().visit(node)\n\n\ndef test_visit_pass_ipython() -> None:\n node = ast.parse(\"pass\").body[0]\n assert isinstance(node, ast.Pass)\n assert (\n algorithmic_codegen.IPythonAlgorithmicCodegen().visit(node) == r\"\\mathbf{pass}\"\n )\n\n\ndef test_visit_break_ipython() -> None:\n node = ast.parse(\"break\").body[0]\n assert isinstance(node, ast.Break)\n assert (\n algorithmic_codegen.IPythonAlgorithmicCodegen().visit(node) == r\"\\mathbf{break}\"\n )\n\n\ndef test_visit_continue_ipython() -> None:\n node = ast.parse(\"continue\").body[0]\n assert isinstance(node, ast.Continue)\n assert (\n algorithmic_codegen.IPythonAlgorithmicCodegen().visit(node)\n == r\"\\mathbf{continue}\"\n )\n"
},
{
"path": "src/latexify/codegen/codegen_utils.py",
"content": "from typing import Any\n\nfrom latexify import exceptions\n\n\ndef convert_constant(value: Any) -> str:\n \"\"\"Helper to convert constant values to LaTeX.\n\n Args:\n value: A constant value.\n\n Returns:\n The LaTeX representation of `value`.\n \"\"\"\n if value is None or isinstance(value, bool):\n return r\"\\mathrm{\" + str(value) + \"}\"\n if isinstance(value, (int, float, complex)):\n # TODO(odashi): Support other symbols for the imaginary unit than j.\n return str(value)\n if isinstance(value, str):\n return r'\\textrm{\"' + value + '\"}'\n if isinstance(value, bytes):\n return r\"\\textrm{\" + str(value) + \"}\"\n if value is ...:\n return r\"\\cdots\"\n raise exceptions.LatexifyNotSupportedError(\n f\"Unrecognized constant: {type(value).__name__}\"\n )\n"
},
{
"path": "src/latexify/codegen/codegen_utils_test.py",
"content": "\"\"\"Tests for latexify.codegen.codegen_utils.\"\"\"\n\nfrom __future__ import annotations\n\nfrom typing import Any\n\nimport pytest\n\nfrom latexify import exceptions\nfrom latexify.codegen.codegen_utils import convert_constant\n\n\n@pytest.mark.parametrize(\n \"constant,latex\",\n [\n (None, r\"\\mathrm{None}\"),\n (True, r\"\\mathrm{True}\"),\n (False, r\"\\mathrm{False}\"),\n (123, \"123\"),\n (456.789, \"456.789\"),\n (-3 + 4j, \"(-3+4j)\"),\n (\"string\", r'\\textrm{\"string\"}'),\n (..., r\"\\cdots\"),\n ],\n)\ndef test_convert_constant(constant: Any, latex: str) -> None:\n assert convert_constant(constant) == latex\n\n\ndef test_convert_constant_unsupported_constant() -> None:\n with pytest.raises(\n exceptions.LatexifyNotSupportedError, match=\"^Unrecognized constant: \"\n ):\n convert_constant({})\n"
},
{
"path": "src/latexify/codegen/expression_codegen.py",
"content": "\"\"\"Codegen for single expressions.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\nimport re\n\nfrom latexify import analyzers, ast_utils, exceptions\nfrom latexify.codegen import codegen_utils, expression_rules, identifier_converter\n\n\nclass ExpressionCodegen(ast.NodeVisitor):\n \"\"\"Codegen for single expressions.\"\"\"\n\n _identifier_converter: identifier_converter.IdentifierConverter\n\n _bin_op_rules: dict[type[ast.operator], expression_rules.BinOpRule]\n _compare_ops: dict[type[ast.cmpop], str]\n\n def __init__(\n self,\n *,\n use_math_symbols: bool = False,\n use_set_symbols: bool = False,\n escape_underscores: bool = True,\n ) -> None:\n \"\"\"Initializer.\n\n Args:\n use_math_symbols: Whether to convert identifiers with a math symbol\n surface (e.g., \"alpha\") to the LaTeX symbol (e.g., \"\\\\alpha\").\n use_set_symbols: Whether to use set symbols or not.\n \"\"\"\n self._identifier_converter = identifier_converter.IdentifierConverter(\n use_math_symbols=use_math_symbols, escape_underscores=escape_underscores\n )\n\n self._bin_op_rules = (\n expression_rules.SET_BIN_OP_RULES\n if use_set_symbols\n else expression_rules.BIN_OP_RULES\n )\n self._compare_ops = (\n expression_rules.SET_COMPARE_OPS\n if use_set_symbols\n else expression_rules.COMPARE_OPS\n )\n\n def generic_visit(self, node: ast.AST) -> str:\n raise exceptions.LatexifyNotSupportedError(\n f\"Unsupported AST: {type(node).__name__}\"\n )\n\n def visit_Tuple(self, node: ast.Tuple) -> str:\n \"\"\"Visit a Tuple node.\"\"\"\n elts = [self.visit(elt) for elt in node.elts]\n return r\"\\mathopen{}\\left( \" + r\", \".join(elts) + r\" \\mathclose{}\\right)\"\n\n def visit_List(self, node: ast.List) -> str:\n \"\"\"Visit a List node.\"\"\"\n elts = [self.visit(elt) for elt in node.elts]\n return r\"\\mathopen{}\\left[ \" + r\", \".join(elts) + r\" \\mathclose{}\\right]\"\n\n def visit_Set(self, node: ast.Set) -> str:\n \"\"\"Visit a Set node.\"\"\"\n elts = [self.visit(elt) for elt in node.elts]\n return r\"\\mathopen{}\\left\\{ \" + r\", \".join(elts) + r\" \\mathclose{}\\right\\}\"\n\n def visit_ListComp(self, node: ast.ListComp) -> str:\n \"\"\"Visit a ListComp node.\"\"\"\n generators = [self.visit(comp) for comp in node.generators]\n return (\n r\"\\mathopen{}\\left[ \"\n + self.visit(node.elt)\n + r\" \\mid \"\n + \", \".join(generators)\n + r\" \\mathclose{}\\right]\"\n )\n\n def visit_SetComp(self, node: ast.SetComp) -> str:\n \"\"\"Visit a SetComp node.\"\"\"\n generators = [self.visit(comp) for comp in node.generators]\n return (\n r\"\\mathopen{}\\left\\{ \"\n + self.visit(node.elt)\n + r\" \\mid \"\n + \", \".join(generators)\n + r\" \\mathclose{}\\right\\}\"\n )\n\n def visit_comprehension(self, node: ast.comprehension) -> str:\n \"\"\"Visit a comprehension node.\"\"\"\n target = rf\"{self.visit(node.target)} \\in {self.visit(node.iter)}\"\n\n if not node.ifs:\n # Returns the source without parenthesis.\n return target\n\n conds = [target] + [self.visit(cond) for cond in node.ifs]\n wrapped = [r\"\\mathopen{}\\left( \" + s + r\" \\mathclose{}\\right)\" for s in conds]\n return r\" \\land \".join(wrapped)\n\n def _generate_sum_prod(self, node: ast.Call) -> str | None:\n \"\"\"Generates sum/prod expression.\n\n Args:\n node: ast.Call node containing the sum/prod invocation.\n\n Returns:\n Generated LaTeX, or None if the node has unsupported syntax.\n \"\"\"\n if not node.args or not isinstance(node.args[0], ast.GeneratorExp):\n return None\n\n name = ast_utils.extract_function_name_or_none(node)\n assert name in (\"fsum\", \"sum\", \"prod\")\n\n command = {\n \"fsum\": r\"\\sum\",\n \"sum\": r\"\\sum\",\n \"prod\": r\"\\prod\",\n }[name]\n\n elt, scripts = self._get_sum_prod_info(node.args[0])\n scripts_str = [rf\"{command}_{{{lo}}}^{{{up}}}\" for lo, up in scripts]\n return (\n \" \".join(scripts_str)\n + rf\" \\mathopen{{}}\\left({{{elt}}}\\mathclose{{}}\\right)\"\n )\n\n def _generate_matrix(self, node: ast.Call) -> str | None:\n \"\"\"Generates matrix expression.\n\n Args:\n node: ast.Call node containing the ndarray invocation.\n\n Returns:\n Generated LaTeX, or None if the node has unsupported syntax.\n \"\"\"\n\n def generate_matrix_from_array(data: list[list[str]]) -> str:\n \"\"\"Helper to generate a bmatrix environment.\"\"\"\n contents = r\" \\\\ \".join(\" & \".join(row) for row in data)\n return r\"\\begin{bmatrix} \" + contents + r\" \\end{bmatrix}\"\n\n arg = node.args[0]\n if not isinstance(arg, ast.List) or not arg.elts:\n # Not an array or no rows\n return None\n\n row0 = arg.elts[0]\n\n if not isinstance(row0, ast.List):\n # Maybe 1 x N array\n return generate_matrix_from_array([[self.visit(x) for x in arg.elts]])\n\n if not row0.elts:\n # No columns\n return None\n\n ncols = len(row0.elts)\n\n rows: list[list[str]] = []\n\n for row in arg.elts:\n if not isinstance(row, ast.List) or len(row.elts) != ncols:\n # Length mismatch\n return None\n\n rows.append([self.visit(x) for x in row.elts])\n\n return generate_matrix_from_array(rows)\n\n def _generate_zeros(self, node: ast.Call) -> str | None:\n \"\"\"Generates LaTeX for numpy.zeros.\n Args:\n node: ast.Call node containing the appropriate method invocation.\n Returns:\n Generated LaTeX, or None if the node has unsupported syntax.\n \"\"\"\n name = ast_utils.extract_function_name_or_none(node)\n assert name == \"zeros\"\n\n if len(node.args) != 1:\n return None\n\n # All args to np.zeros should be numeric.\n if isinstance(node.args[0], ast.Tuple):\n dims = [ast_utils.extract_int_or_none(x) for x in node.args[0].elts]\n if any(x is None for x in dims):\n return None\n if not dims:\n return \"0\"\n if len(dims) == 1:\n dims = [1, dims[0]]\n\n dims_latex = r\" \\times \".join(str(x) for x in dims)\n else:\n dim = ast_utils.extract_int_or_none(node.args[0])\n if not isinstance(dim, int):\n return None\n # 1 x N array of zeros\n dims_latex = rf\"1 \\times {dim}\"\n\n return rf\"\\mathbf{{0}}^{{{dims_latex}}}\"\n\n def _generate_identity(self, node: ast.Call) -> str | None:\n \"\"\"Generates LaTeX for numpy.identity.\n Args:\n node: ast.Call node containing the appropriate method invocation.\n Returns:\n Generated LaTeX, or None if the node has unsupported syntax.\n \"\"\"\n name = ast_utils.extract_function_name_or_none(node)\n assert name == \"identity\"\n\n if len(node.args) != 1:\n return None\n\n ndims = ast_utils.extract_int_or_none(node.args[0])\n if ndims is None:\n return None\n\n return rf\"\\mathbf{{I}}_{{{ndims}}}\"\n\n def _generate_transpose(self, node: ast.Call) -> str | None:\n \"\"\"Generates LaTeX for numpy.transpose.\n Args:\n node: ast.Call node containing the appropriate method invocation.\n Returns:\n Generated LaTeX, or None if the node has unsupported syntax.\n Raises:\n LatexifyError: Unsupported argument type given.\n \"\"\"\n name = ast_utils.extract_function_name_or_none(node)\n assert name == \"transpose\"\n\n if len(node.args) != 1:\n return None\n\n func_arg = node.args[0]\n if isinstance(func_arg, ast.Name):\n return rf\"\\mathbf{{{func_arg.id}}}^\\intercal\"\n else:\n return None\n\n def _generate_determinant(self, node: ast.Call) -> str | None:\n \"\"\"Generates LaTeX for numpy.linalg.det.\n Args:\n node: ast.Call node containing the appropriate method invocation.\n Returns:\n Generated LaTeX, or None if the node has unsupported syntax.\n Raises:\n LatexifyError: Unsupported argument type given.\n \"\"\"\n name = ast_utils.extract_function_name_or_none(node)\n assert name == \"det\"\n\n if len(node.args) != 1:\n return None\n\n func_arg = node.args[0]\n if isinstance(func_arg, ast.Name):\n arg_id = rf\"\\mathbf{{{func_arg.id}}}\"\n return rf\"\\det \\mathopen{{}}\\left( {arg_id} \\mathclose{{}}\\right)\"\n elif isinstance(func_arg, ast.List):\n matrix = self._generate_matrix(node)\n return rf\"\\det \\mathopen{{}}\\left( {matrix} \\mathclose{{}}\\right)\"\n\n return None\n\n def _generate_matrix_rank(self, node: ast.Call) -> str | None:\n \"\"\"Generates LaTeX for numpy.linalg.matrix_rank.\n Args:\n node: ast.Call node containing the appropriate method invocation.\n Returns:\n Generated LaTeX, or None if the node has unsupported syntax.\n Raises:\n LatexifyError: Unsupported argument type given.\n \"\"\"\n name = ast_utils.extract_function_name_or_none(node)\n assert name == \"matrix_rank\"\n\n if len(node.args) != 1:\n return None\n\n func_arg = node.args[0]\n if isinstance(func_arg, ast.Name):\n arg_id = rf\"\\mathbf{{{func_arg.id}}}\"\n return (\n rf\"\\mathrm{{rank}} \\mathopen{{}}\\left( {arg_id} \\mathclose{{}}\\right)\"\n )\n elif isinstance(func_arg, ast.List):\n matrix = self._generate_matrix(node)\n return (\n rf\"\\mathrm{{rank}} \\mathopen{{}}\\left( {matrix} \\mathclose{{}}\\right)\"\n )\n\n return None\n\n def _generate_matrix_power(self, node: ast.Call) -> str | None:\n \"\"\"Generates LaTeX for numpy.linalg.matrix_power.\n Args:\n node: ast.Call node containing the appropriate method invocation.\n Returns:\n Generated LaTeX, or None if the node has unsupported syntax.\n Raises:\n LatexifyError: Unsupported argument type given.\n \"\"\"\n name = ast_utils.extract_function_name_or_none(node)\n assert name == \"matrix_power\"\n\n if len(node.args) != 2:\n return None\n\n func_arg = node.args[0]\n power_arg = node.args[1]\n if isinstance(power_arg, ast.Num):\n if isinstance(func_arg, ast.Name):\n return rf\"\\mathbf{{{func_arg.id}}}^{{{power_arg.n}}}\"\n elif isinstance(func_arg, ast.List):\n matrix = self._generate_matrix(node)\n if matrix is not None:\n return rf\"{matrix}^{{{power_arg.n}}}\"\n return None\n\n def _generate_inv(self, node: ast.Call) -> str | None:\n \"\"\"Generates LaTeX for numpy.linalg.inv.\n Args:\n node: ast.Call node containing the appropriate method invocation.\n Returns:\n Generated LaTeX, or None if the node has unsupported syntax.\n Raises:\n LatexifyError: Unsupported argument type given.\n \"\"\"\n name = ast_utils.extract_function_name_or_none(node)\n assert name == \"inv\"\n\n if len(node.args) != 1:\n return None\n\n func_arg = node.args[0]\n if isinstance(func_arg, ast.Name):\n return rf\"\\mathbf{{{func_arg.id}}}^{{-1}}\"\n elif isinstance(func_arg, ast.List):\n return rf\"{self._generate_matrix(node)}^{{-1}}\"\n return None\n\n def _generate_pinv(self, node: ast.Call) -> str | None:\n \"\"\"Generates LaTeX for numpy.linalg.pinv.\n Args:\n node: ast.Call node containing the appropriate method invocation.\n Returns:\n Generated LaTeX, or None if the node has unsupported syntax.\n Raises:\n LatexifyError: Unsupported argument type given.\n \"\"\"\n name = ast_utils.extract_function_name_or_none(node)\n assert name == \"pinv\"\n\n if len(node.args) != 1:\n return None\n\n func_arg = node.args[0]\n if isinstance(func_arg, ast.Name):\n return rf\"\\mathbf{{{func_arg.id}}}^{{+}}\"\n elif isinstance(func_arg, ast.List):\n return rf\"{self._generate_matrix(node)}^{{+}}\"\n return None\n\n def visit_Call(self, node: ast.Call) -> str:\n \"\"\"Visit a Call node.\"\"\"\n func_name = ast_utils.extract_function_name_or_none(node)\n\n # Special treatments for some functions.\n # TODO(odashi): Move these functions to some separate utility.\n if func_name in (\"fsum\", \"sum\", \"prod\"):\n special_latex = self._generate_sum_prod(node)\n elif func_name in (\"array\", \"ndarray\"):\n special_latex = self._generate_matrix(node)\n elif func_name == \"zeros\":\n special_latex = self._generate_zeros(node)\n elif func_name == \"identity\":\n special_latex = self._generate_identity(node)\n elif func_name == \"transpose\":\n special_latex = self._generate_transpose(node)\n elif func_name == \"det\":\n special_latex = self._generate_determinant(node)\n elif func_name == \"matrix_rank\":\n special_latex = self._generate_matrix_rank(node)\n elif func_name == \"matrix_power\":\n special_latex = self._generate_matrix_power(node)\n elif func_name == \"inv\":\n special_latex = self._generate_inv(node)\n elif func_name == \"pinv\":\n special_latex = self._generate_pinv(node)\n else:\n special_latex = None\n\n if special_latex is not None:\n return special_latex\n\n # Obtains the codegen rule.\n rule = (\n expression_rules.BUILTIN_FUNCS.get(func_name)\n if func_name is not None\n else None\n )\n\n if rule is None:\n rule = expression_rules.FunctionRule(self.visit(node.func))\n\n if rule.is_unary and len(node.args) == 1:\n # Unary function. Applies the same wrapping policy with the unary operators.\n precedence = expression_rules.get_precedence(node)\n arg = node.args[0]\n # NOTE(odashi):\n # Factorial \"x!\" is treated as a special case: it requires both inner/outer\n # parentheses for correct interpretation.\n force_wrap_factorial = isinstance(arg, ast.Call) and (\n func_name == \"factorial\"\n or ast_utils.extract_function_name_or_none(arg) == \"factorial\"\n )\n # Note(odashi):\n # Wrapping is also required if the argument is pow.\n # https://github.com/google/latexify_py/issues/189\n force_wrap_pow = isinstance(arg, ast.BinOp) and isinstance(arg.op, ast.Pow)\n arg_latex = self._wrap_operand(\n arg, precedence, force_wrap_factorial or force_wrap_pow\n )\n elements = [rule.left, arg_latex, rule.right]\n else:\n arg_latex = \", \".join(self.visit(arg) for arg in node.args)\n if rule.is_wrapped:\n elements = [rule.left, arg_latex, rule.right]\n else:\n elements = [\n rule.left,\n r\"\\mathopen{}\\left(\",\n arg_latex,\n r\"\\mathclose{}\\right)\",\n rule.right,\n ]\n\n return \" \".join(x for x in elements if x)\n\n def visit_Attribute(self, node: ast.Attribute) -> str:\n \"\"\"Visit an Attribute node.\"\"\"\n vstr = self.visit(node.value)\n astr = self._identifier_converter.convert(node.attr)[0]\n return vstr + \".\" + astr\n\n def visit_Name(self, node: ast.Name) -> str:\n \"\"\"Visit a Name node.\"\"\"\n return self._identifier_converter.convert(node.id)[0]\n\n # From Python 3.8\n def visit_Constant(self, node: ast.Constant) -> str:\n \"\"\"Visit a Constant node.\"\"\"\n return codegen_utils.convert_constant(node.value)\n\n # Until Python 3.7\n def visit_Num(self, node: ast.Num) -> str:\n \"\"\"Visit a Num node.\"\"\"\n return codegen_utils.convert_constant(node.n)\n\n # Until Python 3.7\n def visit_Str(self, node: ast.Str) -> str:\n \"\"\"Visit a Str node.\"\"\"\n return codegen_utils.convert_constant(node.s)\n\n # Until Python 3.7\n def visit_Bytes(self, node: ast.Bytes) -> str:\n \"\"\"Visit a Bytes node.\"\"\"\n return codegen_utils.convert_constant(node.s)\n\n # Until Python 3.7\n def visit_NameConstant(self, node: ast.NameConstant) -> str:\n \"\"\"Visit a NameConstant node.\"\"\"\n return codegen_utils.convert_constant(node.value)\n\n # Until Python 3.7\n def visit_Ellipsis(self, node: ast.Ellipsis) -> str:\n \"\"\"Visit an Ellipsis node.\"\"\"\n return codegen_utils.convert_constant(...)\n\n def _wrap_operand(\n self, child: ast.expr, parent_prec: int, force_wrap: bool = False\n ) -> str:\n \"\"\"Wraps the operand subtree with parentheses.\n\n Args:\n child: Operand subtree.\n parent_prec: Precedence of the parent operator.\n force_wrap: Whether to wrap the operand or not when the precedence is equal.\n\n Returns:\n LaTeX form of `child`, with or without surrounding parentheses.\n \"\"\"\n latex = self.visit(child)\n child_prec = expression_rules.get_precedence(child)\n\n if force_wrap or child_prec < parent_prec:\n return rf\"\\mathopen{{}}\\left( {latex} \\mathclose{{}}\\right)\"\n\n return latex\n\n def _wrap_binop_operand(\n self,\n child: ast.expr,\n parent_prec: int,\n operand_rule: expression_rules.BinOperandRule,\n ) -> str:\n \"\"\"Wraps the operand subtree of BinOp with parentheses.\n\n Args:\n child: Operand subtree.\n parent_prec: Precedence of the parent operator.\n operand_rule: Syntax rule of this operand.\n\n Returns:\n LaTeX form of the `child`, with or without surrounding parentheses.\n \"\"\"\n if not operand_rule.wrap:\n return self.visit(child)\n\n if isinstance(child, ast.Call):\n child_fn_name = ast_utils.extract_function_name_or_none(child)\n rule = (\n expression_rules.BUILTIN_FUNCS.get(child_fn_name)\n if child_fn_name is not None\n else None\n )\n if rule is not None and rule.is_wrapped:\n return self.visit(child)\n\n if not isinstance(child, ast.BinOp):\n return self._wrap_operand(child, parent_prec)\n\n latex = self.visit(child)\n\n if expression_rules.BIN_OP_RULES[type(child.op)].is_wrapped:\n return latex\n\n child_prec = expression_rules.get_precedence(child)\n\n if child_prec > parent_prec or (\n child_prec == parent_prec and not operand_rule.force\n ):\n return latex\n\n return rf\"\\mathopen{{}}\\left( {latex} \\mathclose{{}}\\right)\"\n\n _l_bracket_pattern = re.compile(r\"^\\\\mathopen.*\")\n _r_bracket_pattern = re.compile(r\".*\\\\mathclose[^ ]+$\")\n _r_word_pattern = re.compile(r\"\\\\mathrm\\{[^ ]+\\}$\")\n\n def _should_remove_multiply_op(\n self, l_latex: str, r_latex: str, l_expr: ast.expr, r_expr: ast.expr\n ):\n \"\"\"Determine whether the multiply operator should be removed or not.\n\n See also:\n https://github.com/google/latexify_py/issues/89#issuecomment-1344967636\n\n This is an ad-hoc implementation.\n This function doesn't fully implements the above requirements, but only\n essential ones necessary to release v0.3.\n \"\"\"\n\n # NOTE(odashi): For compatibility with Python 3.7, we compare the generated\n # caracter type directly to determine the \"numeric\" type.\n\n if isinstance(l_expr, ast.Call):\n l_type = \"f\"\n elif self._r_bracket_pattern.match(l_latex):\n l_type = \"b\"\n elif self._r_word_pattern.match(l_latex):\n l_type = \"w\"\n elif l_latex[-1].isnumeric():\n l_type = \"n\"\n else:\n le = l_expr\n while True:\n if isinstance(le, ast.UnaryOp):\n le = le.operand\n elif isinstance(le, ast.BinOp):\n le = le.right\n elif isinstance(le, ast.Compare):\n le = le.comparators[-1]\n elif isinstance(le, ast.BoolOp):\n le = le.values[-1]\n else:\n break\n l_type = \"a\" if isinstance(le, ast.Name) and len(le.id) == 1 else \"m\"\n\n if isinstance(r_expr, ast.Call):\n r_type = \"f\"\n elif self._l_bracket_pattern.match(r_latex):\n r_type = \"b\"\n elif r_latex.startswith(\"\\\\mathrm\"):\n r_type = \"w\"\n elif r_latex[0].isnumeric():\n r_type = \"n\"\n else:\n re = r_expr\n while True:\n if isinstance(re, ast.UnaryOp):\n if isinstance(re.op, ast.USub):\n # NOTE(odashi): Unary \"-\" always require \\cdot.\n return False\n re = re.operand\n elif isinstance(re, ast.BinOp):\n re = re.left\n elif isinstance(re, ast.Compare):\n re = re.left\n elif isinstance(re, ast.BoolOp):\n re = re.values[0]\n else:\n break\n r_type = \"a\" if isinstance(re, ast.Name) and len(re.id) == 1 else \"m\"\n\n if r_type == \"n\":\n return False\n if l_type in \"bn\":\n return True\n if l_type in \"am\" and r_type in \"am\":\n return True\n return False\n\n def visit_BinOp(self, node: ast.BinOp) -> str:\n \"\"\"Visit a BinOp node.\"\"\"\n prec = expression_rules.get_precedence(node)\n rule = self._bin_op_rules[type(node.op)]\n lhs = self._wrap_binop_operand(node.left, prec, rule.operand_left)\n rhs = self._wrap_binop_operand(node.right, prec, rule.operand_right)\n\n if type(node.op) in [ast.Mult, ast.MatMult]:\n if self._should_remove_multiply_op(lhs, rhs, node.left, node.right):\n return f\"{rule.latex_left}{lhs} {rhs}{rule.latex_right}\"\n\n return f\"{rule.latex_left}{lhs}{rule.latex_middle}{rhs}{rule.latex_right}\"\n\n def visit_UnaryOp(self, node: ast.UnaryOp) -> str:\n \"\"\"Visit a UnaryOp node.\"\"\"\n latex = self._wrap_operand(node.operand, expression_rules.get_precedence(node))\n return expression_rules.UNARY_OPS[type(node.op)] + latex\n\n def visit_Compare(self, node: ast.Compare) -> str:\n \"\"\"Visit a Compare node.\"\"\"\n parent_prec = expression_rules.get_precedence(node)\n lhs = self._wrap_operand(node.left, parent_prec)\n ops = [self._compare_ops[type(x)] for x in node.ops]\n rhs = [self._wrap_operand(x, parent_prec) for x in node.comparators]\n ops_rhs = [f\" {o} {r}\" for o, r in zip(ops, rhs)]\n return lhs + \"\".join(ops_rhs)\n\n def visit_BoolOp(self, node: ast.BoolOp) -> str:\n \"\"\"Visit a BoolOp node.\"\"\"\n parent_prec = expression_rules.get_precedence(node)\n values = [self._wrap_operand(x, parent_prec) for x in node.values]\n op = f\" {expression_rules.BOOL_OPS[type(node.op)]} \"\n return op.join(values)\n\n def visit_IfExp(self, node: ast.IfExp) -> str:\n \"\"\"Visit an IfExp node\"\"\"\n latex = r\"\\left\\{ \\begin{array}{ll} \"\n\n current_expr: ast.expr = node\n\n while isinstance(current_expr, ast.IfExp):\n cond_latex = self.visit(current_expr.test)\n true_latex = self.visit(current_expr.body)\n latex += true_latex + r\", & \\mathrm{if} \\ \" + cond_latex + r\" \\\\ \"\n current_expr = current_expr.orelse\n\n latex += self.visit(current_expr)\n return latex + r\", & \\mathrm{otherwise} \\end{array} \\right.\"\n\n def _get_sum_prod_range(self, node: ast.comprehension) -> tuple[str, str] | None:\n \"\"\"Helper to process range(...) for sum and prod functions.\n\n Args:\n node: comprehension node to be analyzed.\n\n Returns:\n Tuple of following strings:\n - lower_rhs\n - upper\n which are used in _get_sum_prod_info, or None if the analysis failed.\n \"\"\"\n if not (\n isinstance(node.iter, ast.Call)\n and isinstance(node.iter.func, ast.Name)\n and node.iter.func.id == \"range\"\n ):\n return None\n\n try:\n range_info = analyzers.analyze_range(node.iter)\n except exceptions.LatexifyError:\n return None\n\n if (\n # Only accepts ascending order with step size 1.\n range_info.step_int != 1\n or (\n range_info.start_int is not None\n and range_info.stop_int is not None\n and range_info.start_int >= range_info.stop_int\n )\n ):\n return None\n\n if range_info.start_int is None:\n lower_rhs = self.visit(range_info.start)\n else:\n lower_rhs = str(range_info.start_int)\n\n if range_info.stop_int is None:\n upper = self.visit(analyzers.reduce_stop_parameter(range_info.stop))\n else:\n upper = str(range_info.stop_int - 1)\n\n return lower_rhs, upper\n\n def _get_sum_prod_info(\n self, node: ast.GeneratorExp\n ) -> tuple[str, list[tuple[str, str]]]:\n r\"\"\"Process GeneratorExp for sum and prod functions.\n\n Args:\n node: GeneratorExp node to be analyzed.\n\n Returns:\n Tuple of following strings:\n - elt\n - scripts\n which are used to represent sum/prod operators as follows:\n \\sum_{scripts[0][0]}^{scripts[0][1]}\n \\sum_{scripts[1][0]}^{scripts[1][1]}\n ...\n {elt}\n\n Raises:\n LateixfyError: Unsupported AST is given.\n \"\"\"\n elt = self.visit(node.elt)\n\n scripts: list[tuple[str, str]] = []\n\n for comp in node.generators:\n range_args = self._get_sum_prod_range(comp)\n\n if range_args is not None and not comp.ifs:\n target = self.visit(comp.target)\n lower_rhs, upper = range_args\n lower = f\"{target} = {lower_rhs}\"\n else:\n lower = self.visit(comp) # Use a usual comprehension form.\n upper = \"\"\n\n scripts.append((lower, upper))\n\n return elt, scripts\n\n # Until 3.8\n def visit_Index(self, node: ast.Index) -> str:\n \"\"\"Visit an Index node.\"\"\"\n return self.visit(node.value) # type: ignore[attr-defined]\n\n def _convert_nested_subscripts(self, node: ast.Subscript) -> tuple[str, list[str]]:\n \"\"\"Helper function to convert nested subscription.\n\n This function converts x[i][j][...] to \"x\" and [\"i\", \"j\", ...]\n\n Args:\n node: ast.Subscript node to be converted.\n\n Returns:\n Tuple of following strings:\n - The root value of the subscription.\n - Sequence of incices.\n \"\"\"\n if isinstance(node.value, ast.Subscript):\n value, indices = self._convert_nested_subscripts(node.value)\n else:\n value = self.visit(node.value)\n indices = []\n\n indices.append(self.visit(node.slice))\n return value, indices\n\n def visit_Subscript(self, node: ast.Subscript) -> str:\n \"\"\"Visitor a Subscript node.\"\"\"\n value, indices = self._convert_nested_subscripts(node)\n\n # TODO(odashi):\n # \"[i][j][...]\" may be a possible representation as well as \"i, j. ...\"\n indices_str = \", \".join(indices)\n\n return f\"{value}_{{{indices_str}}}\"\n"
},
{
"path": "src/latexify/codegen/expression_codegen_test.py",
"content": "\"\"\"Tests for latexify.codegen.expression_codegen.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\n\nimport pytest\n\nfrom latexify import ast_utils, exceptions\nfrom latexify.codegen import expression_codegen\n\n\ndef test_generic_visit() -> None:\n class UnknownNode(ast.AST):\n pass\n\n with pytest.raises(\n exceptions.LatexifyNotSupportedError,\n match=r\"^Unsupported AST: UnknownNode$\",\n ):\n expression_codegen.ExpressionCodegen().visit(UnknownNode())\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"()\", r\"\\mathopen{}\\left( \\mathclose{}\\right)\"),\n (\"(x,)\", r\"\\mathopen{}\\left( x \\mathclose{}\\right)\"),\n (\"(x, y)\", r\"\\mathopen{}\\left( x, y \\mathclose{}\\right)\"),\n (\"(x, y, z)\", r\"\\mathopen{}\\left( x, y, z \\mathclose{}\\right)\"),\n ],\n)\ndef test_visit_tuple(code: str, latex: str) -> None:\n node = ast_utils.parse_expr(code)\n assert isinstance(node, ast.Tuple)\n assert expression_codegen.ExpressionCodegen().visit(node) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"[]\", r\"\\mathopen{}\\left[ \\mathclose{}\\right]\"),\n (\"[x]\", r\"\\mathopen{}\\left[ x \\mathclose{}\\right]\"),\n (\"[x, y]\", r\"\\mathopen{}\\left[ x, y \\mathclose{}\\right]\"),\n (\"[x, y, z]\", r\"\\mathopen{}\\left[ x, y, z \\mathclose{}\\right]\"),\n ],\n)\ndef test_visit_list(code: str, latex: str) -> None:\n node = ast_utils.parse_expr(code)\n assert isinstance(node, ast.List)\n assert expression_codegen.ExpressionCodegen().visit(node) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n # TODO(odashi): Support set().\n # (\"set()\", r\"\\mathopen{}\\left\\{ \\mathclose{}\\right\\}\"),\n (\"{x}\", r\"\\mathopen{}\\left\\{ x \\mathclose{}\\right\\}\"),\n (\"{x, y}\", r\"\\mathopen{}\\left\\{ x, y \\mathclose{}\\right\\}\"),\n (\"{x, y, z}\", r\"\\mathopen{}\\left\\{ x, y, z \\mathclose{}\\right\\}\"),\n ],\n)\ndef test_visit_set(code: str, latex: str) -> None:\n node = ast_utils.parse_expr(code)\n assert isinstance(node, ast.Set)\n assert expression_codegen.ExpressionCodegen().visit(node) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"[i for i in n]\", r\"\\mathopen{}\\left[ i \\mid i \\in n \\mathclose{}\\right]\"),\n (\n \"[i for i in n if i > 0]\",\n r\"\\mathopen{}\\left[ i \\mid\"\n r\" \\mathopen{}\\left( i \\in n \\mathclose{}\\right)\"\n r\" \\land \\mathopen{}\\left( i > 0 \\mathclose{}\\right)\"\n r\" \\mathclose{}\\right]\",\n ),\n (\n \"[i for i in n if i > 0 if f(i)]\",\n r\"\\mathopen{}\\left[ i \\mid\"\n r\" \\mathopen{}\\left( i \\in n \\mathclose{}\\right)\"\n r\" \\land \\mathopen{}\\left( i > 0 \\mathclose{}\\right)\"\n r\" \\land \\mathopen{}\\left( f \\mathopen{}\\left(\"\n r\" i \\mathclose{}\\right) \\mathclose{}\\right)\"\n r\" \\mathclose{}\\right]\",\n ),\n (\n \"[i for k in n for i in k]\",\n r\"\\mathopen{}\\left[ i \\mid k \\in n, i \\in k\" r\" \\mathclose{}\\right]\",\n ),\n (\n \"[i for k in n for i in k if i > 0]\",\n r\"\\mathopen{}\\left[ i \\mid\"\n r\" k \\in n,\"\n r\" \\mathopen{}\\left( i \\in k \\mathclose{}\\right)\"\n r\" \\land \\mathopen{}\\left( i > 0 \\mathclose{}\\right)\"\n r\" \\mathclose{}\\right]\",\n ),\n (\n \"[i for k in n if f(k) for i in k if i > 0]\",\n r\"\\mathopen{}\\left[ i \\mid\"\n r\" \\mathopen{}\\left( k \\in n \\mathclose{}\\right)\"\n r\" \\land \\mathopen{}\\left( f \\mathopen{}\\left(\"\n r\" k \\mathclose{}\\right) \\mathclose{}\\right),\"\n r\" \\mathopen{}\\left( i \\in k \\mathclose{}\\right)\"\n r\" \\land \\mathopen{}\\left( i > 0 \\mathclose{}\\right)\"\n r\" \\mathclose{}\\right]\",\n ),\n ],\n)\ndef test_visit_listcomp(code: str, latex: str) -> None:\n node = ast_utils.parse_expr(code)\n assert isinstance(node, ast.ListComp)\n assert expression_codegen.ExpressionCodegen().visit(node) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"{i for i in n}\", r\"\\mathopen{}\\left\\{ i \\mid i \\in n \\mathclose{}\\right\\}\"),\n (\n \"{i for i in n if i > 0}\",\n r\"\\mathopen{}\\left\\{ i \\mid\"\n r\" \\mathopen{}\\left( i \\in n \\mathclose{}\\right)\"\n r\" \\land \\mathopen{}\\left( i > 0 \\mathclose{}\\right)\"\n r\" \\mathclose{}\\right\\}\",\n ),\n (\n \"{i for i in n if i > 0 if f(i)}\",\n r\"\\mathopen{}\\left\\{ i \\mid\"\n r\" \\mathopen{}\\left( i \\in n \\mathclose{}\\right)\"\n r\" \\land \\mathopen{}\\left( i > 0 \\mathclose{}\\right)\"\n r\" \\land \\mathopen{}\\left( f \\mathopen{}\\left(\"\n r\" i \\mathclose{}\\right) \\mathclose{}\\right)\"\n r\" \\mathclose{}\\right\\}\",\n ),\n (\n \"{i for k in n for i in k}\",\n r\"\\mathopen{}\\left\\{ i \\mid k \\in n, i \\in k\" r\" \\mathclose{}\\right\\}\",\n ),\n (\n \"{i for k in n for i in k if i > 0}\",\n r\"\\mathopen{}\\left\\{ i \\mid\"\n r\" k \\in n,\"\n r\" \\mathopen{}\\left( i \\in k \\mathclose{}\\right)\"\n r\" \\land \\mathopen{}\\left( i > 0 \\mathclose{}\\right)\"\n r\" \\mathclose{}\\right\\}\",\n ),\n (\n \"{i for k in n if f(k) for i in k if i > 0}\",\n r\"\\mathopen{}\\left\\{ i \\mid\"\n r\" \\mathopen{}\\left( k \\in n \\mathclose{}\\right)\"\n r\" \\land \\mathopen{}\\left( f \\mathopen{}\\left(\"\n r\" k \\mathclose{}\\right) \\mathclose{}\\right),\"\n r\" \\mathopen{}\\left( i \\in k \\mathclose{}\\right)\"\n r\" \\land \\mathopen{}\\left( i > 0 \\mathclose{}\\right)\"\n r\" \\mathclose{}\\right\\}\",\n ),\n ],\n)\ndef test_visit_setcomp(code: str, latex: str) -> None:\n node = ast_utils.parse_expr(code)\n assert isinstance(node, ast.SetComp)\n assert expression_codegen.ExpressionCodegen().visit(node) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"foo(x)\", r\"\\mathrm{foo} \\mathopen{}\\left( x \\mathclose{}\\right)\"),\n (\"f(x)\", r\"f \\mathopen{}\\left( x \\mathclose{}\\right)\"),\n (\"f(-x)\", r\"f \\mathopen{}\\left( -x \\mathclose{}\\right)\"),\n (\"f(x + y)\", r\"f \\mathopen{}\\left( x + y \\mathclose{}\\right)\"),\n (\n \"f(f(x))\",\n r\"f \\mathopen{}\\left(\"\n r\" f \\mathopen{}\\left( x \\mathclose{}\\right)\"\n r\" \\mathclose{}\\right)\",\n ),\n (\"f(sqrt(x))\", r\"f \\mathopen{}\\left( \\sqrt{ x } \\mathclose{}\\right)\"),\n (\"f(sin(x))\", r\"f \\mathopen{}\\left( \\sin x \\mathclose{}\\right)\"),\n (\"f(factorial(x))\", r\"f \\mathopen{}\\left( x ! \\mathclose{}\\right)\"),\n (\"f(x, y)\", r\"f \\mathopen{}\\left( x, y \\mathclose{}\\right)\"),\n (\"sqrt(x)\", r\"\\sqrt{ x }\"),\n (\"sqrt(-x)\", r\"\\sqrt{ -x }\"),\n (\"sqrt(x + y)\", r\"\\sqrt{ x + y }\"),\n (\"sqrt(f(x))\", r\"\\sqrt{ f \\mathopen{}\\left( x \\mathclose{}\\right) }\"),\n (\"sqrt(sqrt(x))\", r\"\\sqrt{ \\sqrt{ x } }\"),\n (\"sqrt(sin(x))\", r\"\\sqrt{ \\sin x }\"),\n (\"sqrt(factorial(x))\", r\"\\sqrt{ x ! }\"),\n (\"sin(x)\", r\"\\sin x\"),\n (\"sin(-x)\", r\"\\sin \\mathopen{}\\left( -x \\mathclose{}\\right)\"),\n (\"sin(x + y)\", r\"\\sin \\mathopen{}\\left( x + y \\mathclose{}\\right)\"),\n (\"sin(f(x))\", r\"\\sin f \\mathopen{}\\left( x \\mathclose{}\\right)\"),\n (\"sin(sqrt(x))\", r\"\\sin \\sqrt{ x }\"),\n (\"sin(sin(x))\", r\"\\sin \\sin x\"),\n (\"sin(factorial(x))\", r\"\\sin \\mathopen{}\\left( x ! \\mathclose{}\\right)\"),\n (\"factorial(x)\", r\"x !\"),\n (\"factorial(-x)\", r\"\\mathopen{}\\left( -x \\mathclose{}\\right) !\"),\n (\"factorial(x + y)\", r\"\\mathopen{}\\left( x + y \\mathclose{}\\right) !\"),\n (\n \"factorial(f(x))\",\n r\"\\mathopen{}\\left(\"\n r\" f \\mathopen{}\\left( x \\mathclose{}\\right)\"\n r\" \\mathclose{}\\right) !\",\n ),\n (\"factorial(sqrt(x))\", r\"\\mathopen{}\\left( \\sqrt{ x } \\mathclose{}\\right) !\"),\n (\"factorial(sin(x))\", r\"\\mathopen{}\\left( \\sin x \\mathclose{}\\right) !\"),\n (\"factorial(factorial(x))\", r\"\\mathopen{}\\left( x ! \\mathclose{}\\right) !\"),\n ],\n)\ndef test_visit_call(code: str, latex: str) -> None:\n node = ast_utils.parse_expr(code)\n assert isinstance(node, ast.Call)\n assert expression_codegen.ExpressionCodegen().visit(node) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"log(x)**2\", r\"\\mathopen{}\\left( \\log x \\mathclose{}\\right)^{2}\"),\n (\"log(x**2)\", r\"\\log \\mathopen{}\\left( x^{2} \\mathclose{}\\right)\"),\n (\n \"log(x**2)**3\",\n r\"\\mathopen{}\\left(\"\n r\" \\log \\mathopen{}\\left( x^{2} \\mathclose{}\\right)\"\n r\" \\mathclose{}\\right)^{3}\",\n ),\n ],\n)\ndef test_visit_call_with_pow(code: str, latex: str) -> None:\n node = ast_utils.parse_expr(code)\n assert isinstance(node, (ast.Call, ast.BinOp))\n assert expression_codegen.ExpressionCodegen().visit(node) == latex\n\n\n@pytest.mark.parametrize(\n \"src_suffix,dest_suffix\",\n [\n # No arguments\n (\"()\", r\" \\mathopen{}\\left( \\mathclose{}\\right)\"),\n # No comprehension\n (\"(x)\", r\" x\"),\n (\n \"([1, 2])\",\n r\" \\mathopen{}\\left[ 1, 2 \\mathclose{}\\right]\",\n ),\n (\n \"({1, 2})\",\n r\" \\mathopen{}\\left\\{ 1, 2 \\mathclose{}\\right\\}\",\n ),\n (\"(f(x))\", r\" f \\mathopen{}\\left( x \\mathclose{}\\right)\"),\n # Single comprehension\n (\"(i for i in x)\", r\"_{i \\in x}^{} \\mathopen{}\\left({i}\\mathclose{}\\right)\"),\n (\n \"(i for i in [1, 2])\",\n r\"_{i \\in \\mathopen{}\\left[ 1, 2 \\mathclose{}\\right]}^{} \"\n r\"\\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in {1, 2})\",\n r\"_{i \\in \\mathopen{}\\left\\{ 1, 2 \\mathclose{}\\right\\}}^{}\"\n r\" \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in f(x))\",\n r\"_{i \\in f \\mathopen{}\\left( x \\mathclose{}\\right)}^{}\"\n r\" \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in range(n))\",\n r\"_{i = 0}^{n - 1} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in range(n + 1))\",\n r\"_{i = 0}^{n} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in range(n + 2))\",\n r\"_{i = 0}^{n + 1} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n # ast.parse() does not recognize negative integers.\n \"(i for i in range(n - -1))\",\n r\"_{i = 0}^{n - -1 - 1} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in range(n - 1))\",\n r\"_{i = 0}^{n - 2} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in range(n + m))\",\n r\"_{i = 0}^{n + m - 1} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in range(n - m))\",\n r\"_{i = 0}^{n - m - 1} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in range(3))\",\n r\"_{i = 0}^{2} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in range(3 + 1))\",\n r\"_{i = 0}^{3} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in range(3 + 2))\",\n r\"_{i = 0}^{3 + 1} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in range(3 - 1))\",\n r\"_{i = 0}^{3 - 2} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n # ast.parse() does not recognize negative integers.\n \"(i for i in range(3 - -1))\",\n r\"_{i = 0}^{3 - -1 - 1} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in range(3 + m))\",\n r\"_{i = 0}^{3 + m - 1} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in range(3 - m))\",\n r\"_{i = 0}^{3 - m - 1} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in range(n, m))\",\n r\"_{i = n}^{m - 1} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in range(1, m))\",\n r\"_{i = 1}^{m - 1} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in range(n, 3))\",\n r\"_{i = n}^{2} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in range(n, m, k))\",\n r\"_{i \\in \\mathrm{range} \\mathopen{}\\left( n, m, k \\mathclose{}\\right)}^{}\"\n r\" \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n ],\n)\ndef test_visit_call_sum_prod(src_suffix: str, dest_suffix: str) -> None:\n for src_fn, dest_fn in [(\"fsum\", r\"\\sum\"), (\"sum\", r\"\\sum\"), (\"prod\", r\"\\prod\")]:\n node = ast_utils.parse_expr(src_fn + src_suffix)\n assert isinstance(node, ast.Call)\n assert (\n expression_codegen.ExpressionCodegen().visit(node) == dest_fn + dest_suffix\n )\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n # 2 clauses\n (\n \"sum(i for y in x for i in y)\",\n r\"\\sum_{y \\in x}^{} \\sum_{i \\in y}^{} \"\n r\"\\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"sum(i for y in x for z in y for i in z)\",\n r\"\\sum_{y \\in x}^{} \\sum_{z \\in y}^{} \\sum_{i \\in z}^{} \"\n r\"\\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n # 3 clauses\n (\n \"prod(i for y in x for i in y)\",\n r\"\\prod_{y \\in x}^{} \\prod_{i \\in y}^{} \"\n r\"\\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"prod(i for y in x for z in y for i in z)\",\n r\"\\prod_{y \\in x}^{} \\prod_{z \\in y}^{} \\prod_{i \\in z}^{} \"\n r\"\\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n # reduce stop parameter\n (\n \"sum(i for i in range(n+1))\",\n r\"\\sum_{i = 0}^{n} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"prod(i for i in range(n-1))\",\n r\"\\prod_{i = 0}^{n - 2} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n # reduce stop parameter\n (\n \"sum(i for i in range(n+1))\",\n r\"\\sum_{i = 0}^{n} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"prod(i for i in range(n-1))\",\n r\"\\prod_{i = 0}^{n - 2} \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n ],\n)\ndef test_visit_call_sum_prod_multiple_comprehension(code: str, latex: str) -> None:\n node = ast_utils.parse_expr(code)\n assert isinstance(node, ast.Call)\n assert expression_codegen.ExpressionCodegen().visit(node) == latex\n\n\n@pytest.mark.parametrize(\n \"src_suffix,dest_suffix\",\n [\n (\n \"(i for i in x if i < y)\",\n r\"_{\\mathopen{}\\left( i \\in x \\mathclose{}\\right) \"\n r\"\\land \\mathopen{}\\left( i < y \\mathclose{}\\right)}^{} \"\n r\"\\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n (\n \"(i for i in x if i < y if f(i))\",\n r\"_{\\mathopen{}\\left( i \\in x \\mathclose{}\\right)\"\n r\" \\land \\mathopen{}\\left( i < y \\mathclose{}\\right)\"\n r\" \\land \\mathopen{}\\left( f \\mathopen{}\\left(\"\n r\" i \\mathclose{}\\right) \\mathclose{}\\right)}^{}\"\n r\" \\mathopen{}\\left({i}\\mathclose{}\\right)\",\n ),\n ],\n)\ndef test_visit_call_sum_prod_with_if(src_suffix: str, dest_suffix: str) -> None:\n for src_fn, dest_fn in [(\"sum\", r\"\\sum\"), (\"prod\", r\"\\prod\")]:\n node = ast_utils.parse_expr(src_fn + src_suffix)\n assert isinstance(node, ast.Call)\n assert (\n expression_codegen.ExpressionCodegen().visit(node) == dest_fn + dest_suffix\n )\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\n \"x if x < y else y\",\n r\"\\left\\{ \\begin{array}{ll}\"\n r\" x, & \\mathrm{if} \\ x < y \\\\\"\n r\" y, & \\mathrm{otherwise}\"\n r\" \\end{array} \\right.\",\n ),\n (\n \"x if x < y else (y if y < z else z)\",\n r\"\\left\\{ \\begin{array}{ll}\"\n r\" x, & \\mathrm{if} \\ x < y \\\\\"\n r\" y, & \\mathrm{if} \\ y < z \\\\\"\n r\" z, & \\mathrm{otherwise}\"\n r\" \\end{array} \\right.\",\n ),\n (\n \"x if x < y else (y if y < z else (z if z < w else w))\",\n r\"\\left\\{ \\begin{array}{ll}\"\n r\" x, & \\mathrm{if} \\ x < y \\\\\"\n r\" y, & \\mathrm{if} \\ y < z \\\\\"\n r\" z, & \\mathrm{if} \\ z < w \\\\\"\n r\" w, & \\mathrm{otherwise}\"\n r\" \\end{array} \\right.\",\n ),\n ],\n)\ndef test_if_then_else(code: str, latex: str) -> None:\n node = ast_utils.parse_expr(code)\n assert isinstance(node, ast.IfExp)\n assert expression_codegen.ExpressionCodegen().visit(node) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n # x op y\n (\"x**y\", r\"x^{y}\"),\n (\"x * y\", r\"x y\"),\n (\"x @ y\", r\"x y\"),\n (\"x / y\", r\"\\frac{x}{y}\"),\n (\"x // y\", r\"\\left\\lfloor\\frac{x}{y}\\right\\rfloor\"),\n (\"x % y\", r\"x \\mathbin{\\%} y\"),\n (\"x + y\", r\"x + y\"),\n (\"x - y\", r\"x - y\"),\n (\"x << y\", r\"x \\ll y\"),\n (\"x >> y\", r\"x \\gg y\"),\n (\"x & y\", r\"x \\mathbin{\\&} y\"),\n (\"x ^ y\", r\"x \\oplus y\"),\n (\"x | y\", R\"x \\mathbin{|} y\"),\n # (x op y) op z\n (\"(x**y)**z\", r\"\\mathopen{}\\left( x^{y} \\mathclose{}\\right)^{z}\"),\n (\"(x * y) * z\", r\"x y z\"),\n (\"(x @ y) @ z\", r\"x y z\"),\n (\"(x / y) / z\", r\"\\frac{\\frac{x}{y}}{z}\"),\n (\n \"(x // y) // z\",\n r\"\\left\\lfloor\\frac{\\left\\lfloor\\frac{x}{y}\\right\\rfloor}{z}\\right\\rfloor\",\n ),\n (\"(x % y) % z\", r\"x \\mathbin{\\%} y \\mathbin{\\%} z\"),\n (\"(x + y) + z\", r\"x + y + z\"),\n (\"(x - y) - z\", r\"x - y - z\"),\n (\"(x << y) << z\", r\"x \\ll y \\ll z\"),\n (\"(x >> y) >> z\", r\"x \\gg y \\gg z\"),\n (\"(x & y) & z\", r\"x \\mathbin{\\&} y \\mathbin{\\&} z\"),\n (\"(x ^ y) ^ z\", r\"x \\oplus y \\oplus z\"),\n (\"(x | y) | z\", r\"x \\mathbin{|} y \\mathbin{|} z\"),\n # x op (y op z)\n (\"x**(y**z)\", r\"x^{y^{z}}\"),\n (\"x * (y * z)\", r\"x y z\"),\n (\"x @ (y @ z)\", r\"x y z\"),\n (\"x / (y / z)\", r\"\\frac{x}{\\frac{y}{z}}\"),\n (\n \"x // (y // z)\",\n r\"\\left\\lfloor\\frac{x}{\\left\\lfloor\\frac{y}{z}\\right\\rfloor}\\right\\rfloor\",\n ),\n (\n \"x % (y % z)\",\n r\"x \\mathbin{\\%} \\mathopen{}\\left( y \\mathbin{\\%} z \\mathclose{}\\right)\",\n ),\n (\"x + (y + z)\", r\"x + y + z\"),\n (\"x - (y - z)\", r\"x - \\mathopen{}\\left( y - z \\mathclose{}\\right)\"),\n (\"x << (y << z)\", r\"x \\ll \\mathopen{}\\left( y \\ll z \\mathclose{}\\right)\"),\n (\"x >> (y >> z)\", r\"x \\gg \\mathopen{}\\left( y \\gg z \\mathclose{}\\right)\"),\n (\"x & (y & z)\", r\"x \\mathbin{\\&} y \\mathbin{\\&} z\"),\n (\"x ^ (y ^ z)\", r\"x \\oplus y \\oplus z\"),\n (\"x | (y | z)\", r\"x \\mathbin{|} y \\mathbin{|} z\"),\n # x OP y op z\n (\"x**y * z\", r\"x^{y} z\"),\n (\"x * y + z\", r\"x y + z\"),\n (\"x @ y + z\", r\"x y + z\"),\n (\"x / y + z\", r\"\\frac{x}{y} + z\"),\n (\"x // y + z\", r\"\\left\\lfloor\\frac{x}{y}\\right\\rfloor + z\"),\n (\"x % y + z\", r\"x \\mathbin{\\%} y + z\"),\n (\"x + y << z\", r\"x + y \\ll z\"),\n (\"x - y << z\", r\"x - y \\ll z\"),\n (\"x << y & z\", r\"x \\ll y \\mathbin{\\&} z\"),\n (\"x >> y & z\", r\"x \\gg y \\mathbin{\\&} z\"),\n (\"x & y ^ z\", r\"x \\mathbin{\\&} y \\oplus z\"),\n (\"x ^ y | z\", r\"x \\oplus y \\mathbin{|} z\"),\n # x OP (y op z)\n (\"x**(y * z)\", r\"x^{y z}\"),\n (\"x * (y + z)\", r\"x \\cdot \\mathopen{}\\left( y + z \\mathclose{}\\right)\"),\n (\"x @ (y + z)\", r\"x \\cdot \\mathopen{}\\left( y + z \\mathclose{}\\right)\"),\n (\"x / (y + z)\", r\"\\frac{x}{y + z}\"),\n (\"x // (y + z)\", r\"\\left\\lfloor\\frac{x}{y + z}\\right\\rfloor\"),\n (\"x % (y + z)\", r\"x \\mathbin{\\%} \\mathopen{}\\left( y + z \\mathclose{}\\right)\"),\n (\"x + (y << z)\", r\"x + \\mathopen{}\\left( y \\ll z \\mathclose{}\\right)\"),\n (\"x - (y << z)\", r\"x - \\mathopen{}\\left( y \\ll z \\mathclose{}\\right)\"),\n (\n \"x << (y & z)\",\n r\"x \\ll \\mathopen{}\\left( y \\mathbin{\\&} z \\mathclose{}\\right)\",\n ),\n (\n \"x >> (y & z)\",\n r\"x \\gg \\mathopen{}\\left( y \\mathbin{\\&} z \\mathclose{}\\right)\",\n ),\n (\n \"x & (y ^ z)\",\n r\"x \\mathbin{\\&} \\mathopen{}\\left( y \\oplus z \\mathclose{}\\right)\",\n ),\n (\n \"x ^ (y | z)\",\n r\"x \\oplus \\mathopen{}\\left( y \\mathbin{|} z \\mathclose{}\\right)\",\n ),\n # x op y OP z\n (\"x * y**z\", r\"x y^{z}\"),\n (\"x + y * z\", r\"x + y z\"),\n (\"x + y @ z\", r\"x + y z\"),\n (\"x + y / z\", r\"x + \\frac{y}{z}\"),\n (\"x + y // z\", r\"x + \\left\\lfloor\\frac{y}{z}\\right\\rfloor\"),\n (\"x + y % z\", r\"x + y \\mathbin{\\%} z\"),\n (\"x << y + z\", r\"x \\ll y + z\"),\n (\"x << y - z\", r\"x \\ll y - z\"),\n (\"x & y << z\", r\"x \\mathbin{\\&} y \\ll z\"),\n (\"x & y >> z\", r\"x \\mathbin{\\&} y \\gg z\"),\n (\"x ^ y & z\", r\"x \\oplus y \\mathbin{\\&} z\"),\n (\"x | y ^ z\", r\"x \\mathbin{|} y \\oplus z\"),\n # (x op y) OP z\n (\"(x * y)**z\", r\"\\mathopen{}\\left( x y \\mathclose{}\\right)^{z}\"),\n (\"(x + y) * z\", r\"\\mathopen{}\\left( x + y \\mathclose{}\\right) z\"),\n (\"(x + y) @ z\", r\"\\mathopen{}\\left( x + y \\mathclose{}\\right) z\"),\n (\"(x + y) / z\", r\"\\frac{x + y}{z}\"),\n (\"(x + y) // z\", r\"\\left\\lfloor\\frac{x + y}{z}\\right\\rfloor\"),\n (\"(x + y) % z\", r\"\\mathopen{}\\left( x + y \\mathclose{}\\right) \\mathbin{\\%} z\"),\n (\"(x << y) + z\", r\"\\mathopen{}\\left( x \\ll y \\mathclose{}\\right) + z\"),\n (\"(x << y) - z\", r\"\\mathopen{}\\left( x \\ll y \\mathclose{}\\right) - z\"),\n (\n \"(x & y) << z\",\n r\"\\mathopen{}\\left( x \\mathbin{\\&} y \\mathclose{}\\right) \\ll z\",\n ),\n (\n \"(x & y) >> z\",\n r\"\\mathopen{}\\left( x \\mathbin{\\&} y \\mathclose{}\\right) \\gg z\",\n ),\n (\n \"(x ^ y) & z\",\n r\"\\mathopen{}\\left( x \\oplus y \\mathclose{}\\right) \\mathbin{\\&} z\",\n ),\n (\n \"(x | y) ^ z\",\n r\"\\mathopen{}\\left( x \\mathbin{|} y \\mathclose{}\\right) \\oplus z\",\n ),\n # is_wrapped\n (\"(x // y)**z\", r\"\\left\\lfloor\\frac{x}{y}\\right\\rfloor^{z}\"),\n # With Call\n (\"x**f(y)\", r\"x^{f \\mathopen{}\\left( y \\mathclose{}\\right)}\"),\n (\n \"f(x)**y\",\n r\"\\mathopen{}\\left(\"\n r\" f \\mathopen{}\\left( x \\mathclose{}\\right)\"\n r\" \\mathclose{}\\right)^{y}\",\n ),\n (\"x * f(y)\", r\"x \\cdot f \\mathopen{}\\left( y \\mathclose{}\\right)\"),\n (\"f(x) * y\", r\"f \\mathopen{}\\left( x \\mathclose{}\\right) \\cdot y\"),\n (\"x / f(y)\", r\"\\frac{x}{f \\mathopen{}\\left( y \\mathclose{}\\right)}\"),\n (\"f(x) / y\", r\"\\frac{f \\mathopen{}\\left( x \\mathclose{}\\right)}{y}\"),\n (\"x + f(y)\", r\"x + f \\mathopen{}\\left( y \\mathclose{}\\right)\"),\n (\"f(x) + y\", r\"f \\mathopen{}\\left( x \\mathclose{}\\right) + y\"),\n # With is_wrapped Call\n (\"sqrt(x) ** y\", r\"\\sqrt{ x }^{y}\"),\n # With UnaryOp\n (\"x**-y\", r\"x^{-y}\"),\n (\"(-x)**y\", r\"\\mathopen{}\\left( -x \\mathclose{}\\right)^{y}\"),\n (\"x * -y\", r\"x \\cdot -y\"),\n (\"-x * y\", r\"-x y\"),\n (\"x / -y\", r\"\\frac{x}{-y}\"),\n (\"-x / y\", r\"\\frac{-x}{y}\"),\n (\"x + -y\", r\"x + -y\"),\n (\"-x + y\", r\"-x + y\"),\n # With Compare\n (\"x**(y == z)\", r\"x^{y = z}\"),\n (\"(x == y)**z\", r\"\\mathopen{}\\left( x = y \\mathclose{}\\right)^{z}\"),\n (\"x * (y == z)\", r\"x \\cdot \\mathopen{}\\left( y = z \\mathclose{}\\right)\"),\n (\"(x == y) * z\", r\"\\mathopen{}\\left( x = y \\mathclose{}\\right) z\"),\n (\"x / (y == z)\", r\"\\frac{x}{y = z}\"),\n (\"(x == y) / z\", r\"\\frac{x = y}{z}\"),\n (\"x + (y == z)\", r\"x + \\mathopen{}\\left( y = z \\mathclose{}\\right)\"),\n (\"(x == y) + z\", r\"\\mathopen{}\\left( x = y \\mathclose{}\\right) + z\"),\n # With BoolOp\n (\"x**(y and z)\", r\"x^{y \\land z}\"),\n (\"(x and y)**z\", r\"\\mathopen{}\\left( x \\land y \\mathclose{}\\right)^{z}\"),\n (\"x * (y and z)\", r\"x \\cdot \\mathopen{}\\left( y \\land z \\mathclose{}\\right)\"),\n (\"(x and y) * z\", r\"\\mathopen{}\\left( x \\land y \\mathclose{}\\right) z\"),\n (\"x / (y and z)\", r\"\\frac{x}{y \\land z}\"),\n (\"(x and y) / z\", r\"\\frac{x \\land y}{z}\"),\n (\"x + (y and z)\", r\"x + \\mathopen{}\\left( y \\land z \\mathclose{}\\right)\"),\n (\"(x and y) + z\", r\"\\mathopen{}\\left( x \\land y \\mathclose{}\\right) + z\"),\n ],\n)\ndef test_visit_binop(code: str, latex: str) -> None:\n tree = ast_utils.parse_expr(code)\n assert isinstance(tree, ast.BinOp)\n assert expression_codegen.ExpressionCodegen().visit(tree) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n # With literals\n (\"+x\", r\"+x\"),\n (\"-x\", r\"-x\"),\n (\"~x\", r\"\\mathord{\\sim} x\"),\n (\"not x\", r\"\\lnot x\"),\n # With Call\n (\"+f(x)\", r\"+f \\mathopen{}\\left( x \\mathclose{}\\right)\"),\n (\"-f(x)\", r\"-f \\mathopen{}\\left( x \\mathclose{}\\right)\"),\n (\"~f(x)\", r\"\\mathord{\\sim} f \\mathopen{}\\left( x \\mathclose{}\\right)\"),\n (\"not f(x)\", r\"\\lnot f \\mathopen{}\\left( x \\mathclose{}\\right)\"),\n # With BinOp\n (\"+(x + y)\", r\"+\\mathopen{}\\left( x + y \\mathclose{}\\right)\"),\n (\"-(x + y)\", r\"-\\mathopen{}\\left( x + y \\mathclose{}\\right)\"),\n (\"~(x + y)\", r\"\\mathord{\\sim} \\mathopen{}\\left( x + y \\mathclose{}\\right)\"),\n (\"not x + y\", r\"\\lnot \\mathopen{}\\left( x + y \\mathclose{}\\right)\"),\n # With Compare\n (\"+(x == y)\", r\"+\\mathopen{}\\left( x = y \\mathclose{}\\right)\"),\n (\"-(x == y)\", r\"-\\mathopen{}\\left( x = y \\mathclose{}\\right)\"),\n (\"~(x == y)\", r\"\\mathord{\\sim} \\mathopen{}\\left( x = y \\mathclose{}\\right)\"),\n (\"not x == y\", r\"\\lnot \\mathopen{}\\left( x = y \\mathclose{}\\right)\"),\n # With BoolOp\n (\"+(x and y)\", r\"+\\mathopen{}\\left( x \\land y \\mathclose{}\\right)\"),\n (\"-(x and y)\", r\"-\\mathopen{}\\left( x \\land y \\mathclose{}\\right)\"),\n (\n \"~(x and y)\",\n r\"\\mathord{\\sim} \\mathopen{}\\left( x \\land y \\mathclose{}\\right)\",\n ),\n (\"not (x and y)\", r\"\\lnot \\mathopen{}\\left( x \\land y \\mathclose{}\\right)\"),\n ],\n)\ndef test_visit_unaryop(code: str, latex: str) -> None:\n tree = ast_utils.parse_expr(code)\n assert isinstance(tree, ast.UnaryOp)\n assert expression_codegen.ExpressionCodegen().visit(tree) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n # 1 comparator\n (\"a == b\", \"a = b\"),\n (\"a > b\", \"a > b\"),\n (\"a >= b\", r\"a \\ge b\"),\n (\"a in b\", r\"a \\in b\"),\n (\"a is b\", r\"a \\equiv b\"),\n (\"a is not b\", r\"a \\not\\equiv b\"),\n (\"a < b\", \"a < b\"),\n (\"a <= b\", r\"a \\le b\"),\n (\"a != b\", r\"a \\ne b\"),\n (\"a not in b\", r\"a \\notin b\"),\n # 2 comparators\n (\"a == b == c\", \"a = b = c\"),\n (\"a == b > c\", \"a = b > c\"),\n (\"a == b >= c\", r\"a = b \\ge c\"),\n (\"a == b < c\", \"a = b < c\"),\n (\"a == b <= c\", r\"a = b \\le c\"),\n (\"a > b == c\", \"a > b = c\"),\n (\"a > b > c\", \"a > b > c\"),\n (\"a > b >= c\", r\"a > b \\ge c\"),\n (\"a >= b == c\", r\"a \\ge b = c\"),\n (\"a >= b > c\", r\"a \\ge b > c\"),\n (\"a >= b >= c\", r\"a \\ge b \\ge c\"),\n (\"a < b == c\", \"a < b = c\"),\n (\"a < b < c\", \"a < b < c\"),\n (\"a < b <= c\", r\"a < b \\le c\"),\n (\"a <= b == c\", r\"a \\le b = c\"),\n (\"a <= b < c\", r\"a \\le b < c\"),\n (\"a <= b <= c\", r\"a \\le b \\le c\"),\n # With Call\n (\"a == f(b)\", r\"a = f \\mathopen{}\\left( b \\mathclose{}\\right)\"),\n (\"f(a) == b\", r\"f \\mathopen{}\\left( a \\mathclose{}\\right) = b\"),\n # With BinOp\n (\"a == b + c\", r\"a = b + c\"),\n (\"a + b == c\", r\"a + b = c\"),\n # With UnaryOp\n (\"a == -b\", r\"a = -b\"),\n (\"-a == b\", r\"-a = b\"),\n (\"a == (not b)\", r\"a = \\lnot b\"),\n (\"(not a) == b\", r\"\\lnot a = b\"),\n # With BoolOp\n (\"a == (b and c)\", r\"a = \\mathopen{}\\left( b \\land c \\mathclose{}\\right)\"),\n (\"(a and b) == c\", r\"\\mathopen{}\\left( a \\land b \\mathclose{}\\right) = c\"),\n ],\n)\ndef test_visit_compare(code: str, latex: str) -> None:\n tree = ast_utils.parse_expr(code)\n assert isinstance(tree, ast.Compare)\n assert expression_codegen.ExpressionCodegen().visit(tree) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n # With literals\n (\"a and b\", r\"a \\land b\"),\n (\"a and b and c\", r\"a \\land b \\land c\"),\n (\"a or b\", r\"a \\lor b\"),\n (\"a or b or c\", r\"a \\lor b \\lor c\"),\n (\"a or b and c\", r\"a \\lor b \\land c\"),\n (\n \"(a or b) and c\",\n r\"\\mathopen{}\\left( a \\lor b \\mathclose{}\\right) \\land c\",\n ),\n (\"a and b or c\", r\"a \\land b \\lor c\"),\n (\n \"a and (b or c)\",\n r\"a \\land \\mathopen{}\\left( b \\lor c \\mathclose{}\\right)\",\n ),\n # With Call\n (\"a and f(b)\", r\"a \\land f \\mathopen{}\\left( b \\mathclose{}\\right)\"),\n (\"f(a) and b\", r\"f \\mathopen{}\\left( a \\mathclose{}\\right) \\land b\"),\n (\"a or f(b)\", r\"a \\lor f \\mathopen{}\\left( b \\mathclose{}\\right)\"),\n (\"f(a) or b\", r\"f \\mathopen{}\\left( a \\mathclose{}\\right) \\lor b\"),\n # With BinOp\n (\"a and b + c\", r\"a \\land b + c\"),\n (\"a + b and c\", r\"a + b \\land c\"),\n (\"a or b + c\", r\"a \\lor b + c\"),\n (\"a + b or c\", r\"a + b \\lor c\"),\n # With UnaryOp\n (\"a and not b\", r\"a \\land \\lnot b\"),\n (\"not a and b\", r\"\\lnot a \\land b\"),\n (\"a or not b\", r\"a \\lor \\lnot b\"),\n (\"not a or b\", r\"\\lnot a \\lor b\"),\n # With Compare\n (\"a and b == c\", r\"a \\land b = c\"),\n (\"a == b and c\", r\"a = b \\land c\"),\n (\"a or b == c\", r\"a \\lor b = c\"),\n (\"a == b or c\", r\"a = b \\lor c\"),\n ],\n)\ndef test_visit_boolop(code: str, latex: str) -> None:\n tree = ast_utils.parse_expr(code)\n assert isinstance(tree, ast.BoolOp)\n assert expression_codegen.ExpressionCodegen().visit(tree) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"0\", \"0\"),\n (\"1\", \"1\"),\n (\"0.0\", \"0.0\"),\n (\"1.5\", \"1.5\"),\n (\"0.0j\", \"0j\"),\n (\"1.0j\", \"1j\"),\n (\"1.5j\", \"1.5j\"),\n ('\"abc\"', r'\\textrm{\"abc\"}'),\n ('b\"abc\"', r\"\\textrm{b'abc'}\"),\n (\"None\", r\"\\mathrm{None}\"),\n (\"False\", r\"\\mathrm{False}\"),\n (\"True\", r\"\\mathrm{True}\"),\n (\"...\", r\"\\cdots\"),\n ],\n)\ndef test_visit_constant(code: str, latex: str) -> None:\n tree = ast_utils.parse_expr(code)\n assert isinstance(tree, ast.Constant)\n assert expression_codegen.ExpressionCodegen().visit(tree) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"x[0]\", \"x_{0}\"),\n (\"x[0][1]\", \"x_{0, 1}\"),\n (\"x[0][1][2]\", \"x_{0, 1, 2}\"),\n (\"x[foo]\", r\"x_{\\mathrm{foo}}\"),\n (\"x[floor(x)]\", r\"x_{\\mathopen{}\\left\\lfloor x \\mathclose{}\\right\\rfloor}\"),\n ],\n)\ndef test_visit_subscript(code: str, latex: str) -> None:\n tree = ast_utils.parse_expr(code)\n assert isinstance(tree, ast.Subscript)\n assert expression_codegen.ExpressionCodegen().visit(tree) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"a - b\", r\"a \\setminus b\"),\n (\"a & b\", r\"a \\cap b\"),\n (\"a ^ b\", r\"a \\mathbin{\\triangle} b\"),\n (\"a | b\", r\"a \\cup b\"),\n ],\n)\ndef test_visit_binop_use_set_symbols(code: str, latex: str) -> None:\n tree = ast_utils.parse_expr(code)\n assert isinstance(tree, ast.BinOp)\n assert (\n expression_codegen.ExpressionCodegen(use_set_symbols=True).visit(tree) == latex\n )\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"a < b\", r\"a \\subset b\"),\n (\"a <= b\", r\"a \\subseteq b\"),\n (\"a > b\", r\"a \\supset b\"),\n (\"a >= b\", r\"a \\supseteq b\"),\n ],\n)\ndef test_visit_compare_use_set_symbols(code: str, latex: str) -> None:\n tree = ast_utils.parse_expr(code)\n assert isinstance(tree, ast.Compare)\n assert (\n expression_codegen.ExpressionCodegen(use_set_symbols=True).visit(tree) == latex\n )\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"array(1)\", r\"\\mathrm{array} \\mathopen{}\\left( 1 \\mathclose{}\\right)\"),\n (\n \"array([])\",\n r\"\\mathrm{array} \\mathopen{}\\left(\"\n r\" \\mathopen{}\\left[ \\mathclose{}\\right]\"\n r\" \\mathclose{}\\right)\",\n ),\n (\"array([1])\", r\"\\begin{bmatrix} 1 \\end{bmatrix}\"),\n (\"array([1, 2, 3])\", r\"\\begin{bmatrix} 1 & 2 & 3 \\end{bmatrix}\"),\n (\n \"array([[]])\",\n r\"\\mathrm{array} \\mathopen{}\\left(\"\n r\" \\mathopen{}\\left[ \\mathopen{}\\left[\"\n r\" \\mathclose{}\\right] \\mathclose{}\\right]\"\n r\" \\mathclose{}\\right)\",\n ),\n (\"array([[1]])\", r\"\\begin{bmatrix} 1 \\end{bmatrix}\"),\n (\"array([[1], [2], [3]])\", r\"\\begin{bmatrix} 1 \\\\ 2 \\\\ 3 \\end{bmatrix}\"),\n (\n \"array([[1], [2], [3, 4]])\",\n r\"\\mathrm{array} \\mathopen{}\\left(\"\n r\" \\mathopen{}\\left[\"\n r\" \\mathopen{}\\left[ 1 \\mathclose{}\\right],\"\n r\" \\mathopen{}\\left[ 2 \\mathclose{}\\right],\"\n r\" \\mathopen{}\\left[ 3, 4 \\mathclose{}\\right]\"\n r\" \\mathclose{}\\right]\"\n r\" \\mathclose{}\\right)\",\n ),\n (\n \"array([[1, 2], [3, 4], [5, 6]])\",\n r\"\\begin{bmatrix} 1 & 2 \\\\ 3 & 4 \\\\ 5 & 6 \\end{bmatrix}\",\n ),\n # Only checks two cases for ndarray.\n (\"ndarray(1)\", r\"\\mathrm{ndarray} \\mathopen{}\\left( 1 \\mathclose{}\\right)\"),\n (\"ndarray([1])\", r\"\\begin{bmatrix} 1 \\end{bmatrix}\"),\n ],\n)\ndef test_numpy_array(code: str, latex: str) -> None:\n tree = ast_utils.parse_expr(code)\n assert isinstance(tree, ast.Call)\n assert expression_codegen.ExpressionCodegen().visit(tree) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"zeros(0)\", r\"\\mathbf{0}^{1 \\times 0}\"),\n (\"zeros(1)\", r\"\\mathbf{0}^{1 \\times 1}\"),\n (\"zeros(2)\", r\"\\mathbf{0}^{1 \\times 2}\"),\n (\"zeros(())\", r\"0\"),\n (\"zeros((0,))\", r\"\\mathbf{0}^{1 \\times 0}\"),\n (\"zeros((1,))\", r\"\\mathbf{0}^{1 \\times 1}\"),\n (\"zeros((2,))\", r\"\\mathbf{0}^{1 \\times 2}\"),\n (\"zeros((0, 0))\", r\"\\mathbf{0}^{0 \\times 0}\"),\n (\"zeros((1, 1))\", r\"\\mathbf{0}^{1 \\times 1}\"),\n (\"zeros((2, 3))\", r\"\\mathbf{0}^{2 \\times 3}\"),\n (\"zeros((0, 0, 0))\", r\"\\mathbf{0}^{0 \\times 0 \\times 0}\"),\n (\"zeros((1, 1, 1))\", r\"\\mathbf{0}^{1 \\times 1 \\times 1}\"),\n (\"zeros((2, 3, 5))\", r\"\\mathbf{0}^{2 \\times 3 \\times 5}\"),\n # Unsupported\n (\"zeros()\", r\"\\mathrm{zeros} \\mathopen{}\\left( \\mathclose{}\\right)\"),\n (\"zeros(x)\", r\"\\mathrm{zeros} \\mathopen{}\\left( x \\mathclose{}\\right)\"),\n (\"zeros(0, x)\", r\"\\mathrm{zeros} \\mathopen{}\\left( 0, x \\mathclose{}\\right)\"),\n (\n \"zeros((x,))\",\n r\"\\mathrm{zeros} \\mathopen{}\\left(\"\n r\" \\mathopen{}\\left( x \\mathclose{}\\right)\"\n r\" \\mathclose{}\\right)\",\n ),\n ],\n)\ndef test_zeros(code: str, latex: str) -> None:\n tree = ast_utils.parse_expr(code)\n assert isinstance(tree, ast.Call)\n assert expression_codegen.ExpressionCodegen().visit(tree) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"identity(0)\", r\"\\mathbf{I}_{0}\"),\n (\"identity(1)\", r\"\\mathbf{I}_{1}\"),\n (\"identity(2)\", r\"\\mathbf{I}_{2}\"),\n # Unsupported\n (\"identity()\", r\"\\mathrm{identity} \\mathopen{}\\left( \\mathclose{}\\right)\"),\n (\"identity(x)\", r\"\\mathrm{identity} \\mathopen{}\\left( x \\mathclose{}\\right)\"),\n (\n \"identity(0, x)\",\n r\"\\mathrm{identity} \\mathopen{}\\left( 0, x \\mathclose{}\\right)\",\n ),\n ],\n)\ndef test_identity(code: str, latex: str) -> None:\n tree = ast_utils.parse_expr(code)\n assert isinstance(tree, ast.Call)\n assert expression_codegen.ExpressionCodegen().visit(tree) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"transpose(A)\", r\"\\mathbf{A}^\\intercal\"),\n (\"transpose(b)\", r\"\\mathbf{b}^\\intercal\"),\n # Unsupported\n (\"transpose()\", r\"\\mathrm{transpose} \\mathopen{}\\left( \\mathclose{}\\right)\"),\n (\"transpose(2)\", r\"\\mathrm{transpose} \\mathopen{}\\left( 2 \\mathclose{}\\right)\"),\n (\n \"transpose(a, (1, 0))\",\n r\"\\mathrm{transpose} \\mathopen{}\\left( a, \"\n r\"\\mathopen{}\\left( 1, 0 \\mathclose{}\\right) \\mathclose{}\\right)\",\n ),\n ],\n)\ndef test_transpose(code: str, latex: str) -> None:\n tree = ast_utils.parse_expr(code)\n assert isinstance(tree, ast.Call)\n assert expression_codegen.ExpressionCodegen().visit(tree) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"det(A)\", r\"\\det \\mathopen{}\\left( \\mathbf{A} \\mathclose{}\\right)\"),\n (\"det(b)\", r\"\\det \\mathopen{}\\left( \\mathbf{b} \\mathclose{}\\right)\"),\n (\n \"det([[1, 2], [3, 4]])\",\n r\"\\det \\mathopen{}\\left( \\begin{bmatrix} 1 & 2 \\\\\"\n r\" 3 & 4 \\end{bmatrix} \\mathclose{}\\right)\",\n ),\n (\n \"det([[1, 2, 3], [4, 5, 6], [7, 8, 9]])\",\n r\"\\det \\mathopen{}\\left( \\begin{bmatrix} 1 & 2 & 3 \\\\ 4 & 5 & 6 \\\\\"\n r\" 7 & 8 & 9 \\end{bmatrix} \\mathclose{}\\right)\",\n ),\n # Unsupported\n (\"det()\", r\"\\mathrm{det} \\mathopen{}\\left( \\mathclose{}\\right)\"),\n (\"det(2)\", r\"\\mathrm{det} \\mathopen{}\\left( 2 \\mathclose{}\\right)\"),\n (\n \"det(a, (1, 0))\",\n r\"\\mathrm{det} \\mathopen{}\\left( a, \"\n r\"\\mathopen{}\\left( 1, 0 \\mathclose{}\\right) \\mathclose{}\\right)\",\n ),\n ],\n)\ndef test_determinant(code: str, latex: str) -> None:\n tree = ast_utils.parse_expr(code)\n assert isinstance(tree, ast.Call)\n assert expression_codegen.ExpressionCodegen().visit(tree) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\n \"matrix_rank(A)\",\n r\"\\mathrm{rank} \\mathopen{}\\left( \\mathbf{A} \\mathclose{}\\right)\",\n ),\n (\n \"matrix_rank(b)\",\n r\"\\mathrm{rank} \\mathopen{}\\left( \\mathbf{b} \\mathclose{}\\right)\",\n ),\n (\n \"matrix_rank([[1, 2], [3, 4]])\",\n r\"\\mathrm{rank} \\mathopen{}\\left( \\begin{bmatrix} 1 & 2 \\\\\"\n r\" 3 & 4 \\end{bmatrix} \\mathclose{}\\right)\",\n ),\n (\n \"matrix_rank([[1, 2, 3], [4, 5, 6], [7, 8, 9]])\",\n r\"\\mathrm{rank} \\mathopen{}\\left( \\begin{bmatrix} 1 & 2 & 3 \\\\ 4 & 5 & 6 \\\\\"\n r\" 7 & 8 & 9 \\end{bmatrix} \\mathclose{}\\right)\",\n ),\n # Unsupported\n (\n \"matrix_rank()\",\n r\"\\mathrm{matrix\\_rank} \\mathopen{}\\left( \\mathclose{}\\right)\",\n ),\n (\n \"matrix_rank(2)\",\n r\"\\mathrm{matrix\\_rank} \\mathopen{}\\left( 2 \\mathclose{}\\right)\",\n ),\n (\n \"matrix_rank(a, (1, 0))\",\n r\"\\mathrm{matrix\\_rank} \\mathopen{}\\left( a, \"\n r\"\\mathopen{}\\left( 1, 0 \\mathclose{}\\right) \\mathclose{}\\right)\",\n ),\n ],\n)\ndef test_matrix_rank(code: str, latex: str) -> None:\n tree = ast_utils.parse_expr(code)\n assert isinstance(tree, ast.Call)\n assert expression_codegen.ExpressionCodegen().visit(tree) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"matrix_power(A, 2)\", r\"\\mathbf{A}^{2}\"),\n (\"matrix_power(b, 2)\", r\"\\mathbf{b}^{2}\"),\n (\n \"matrix_power([[1, 2], [3, 4]], 2)\",\n r\"\\begin{bmatrix} 1 & 2 \\\\ 3 & 4 \\end{bmatrix}^{2}\",\n ),\n (\n \"matrix_power([[1, 2, 3], [4, 5, 6], [7, 8, 9]], 42)\",\n r\"\\begin{bmatrix} 1 & 2 & 3 \\\\ 4 & 5 & 6 \\\\ 7 & 8 & 9 \\end{bmatrix}^{42}\",\n ),\n # Unsupported\n (\n \"matrix_power()\",\n r\"\\mathrm{matrix\\_power} \\mathopen{}\\left( \\mathclose{}\\right)\",\n ),\n (\n \"matrix_power(2)\",\n r\"\\mathrm{matrix\\_power} \\mathopen{}\\left( 2 \\mathclose{}\\right)\",\n ),\n (\n \"matrix_power(a, (1, 0))\",\n r\"\\mathrm{matrix\\_power} \\mathopen{}\\left( a, \"\n r\"\\mathopen{}\\left( 1, 0 \\mathclose{}\\right) \\mathclose{}\\right)\",\n ),\n ],\n)\ndef test_matrix_power(code: str, latex: str) -> None:\n tree = ast_utils.parse_expr(code)\n assert isinstance(tree, ast.Call)\n assert expression_codegen.ExpressionCodegen().visit(tree) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"inv(A)\", r\"\\mathbf{A}^{-1}\"),\n (\"inv(b)\", r\"\\mathbf{b}^{-1}\"),\n (\"inv([[1, 2], [3, 4]])\", r\"\\begin{bmatrix} 1 & 2 \\\\ 3 & 4 \\end{bmatrix}^{-1}\"),\n (\n \"inv([[1, 2, 3], [4, 5, 6], [7, 8, 9]])\",\n r\"\\begin{bmatrix} 1 & 2 & 3 \\\\ 4 & 5 & 6 \\\\ 7 & 8 & 9 \\end{bmatrix}^{-1}\",\n ),\n # Unsupported\n (\"inv()\", r\"\\mathrm{inv} \\mathopen{}\\left( \\mathclose{}\\right)\"),\n (\"inv(2)\", r\"\\mathrm{inv} \\mathopen{}\\left( 2 \\mathclose{}\\right)\"),\n (\n \"inv(a, (1, 0))\",\n r\"\\mathrm{inv} \\mathopen{}\\left( a, \"\n r\"\\mathopen{}\\left( 1, 0 \\mathclose{}\\right) \\mathclose{}\\right)\",\n ),\n ],\n)\ndef test_inv(code: str, latex: str) -> None:\n tree = ast_utils.parse_expr(code)\n assert isinstance(tree, ast.Call)\n assert expression_codegen.ExpressionCodegen().visit(tree) == latex\n\n\n@pytest.mark.parametrize(\n \"code,latex\",\n [\n (\"pinv(A)\", r\"\\mathbf{A}^{+}\"),\n (\"pinv(b)\", r\"\\mathbf{b}^{+}\"),\n (\"pinv([[1, 2], [3, 4]])\", r\"\\begin{bmatrix} 1 & 2 \\\\ 3 & 4 \\end{bmatrix}^{+}\"),\n (\n \"pinv([[1, 2, 3], [4, 5, 6], [7, 8, 9]])\",\n r\"\\begin{bmatrix} 1 & 2 & 3 \\\\ 4 & 5 & 6 \\\\ 7 & 8 & 9 \\end{bmatrix}^{+}\",\n ),\n # Unsupported\n (\"pinv()\", r\"\\mathrm{pinv} \\mathopen{}\\left( \\mathclose{}\\right)\"),\n (\"pinv(2)\", r\"\\mathrm{pinv} \\mathopen{}\\left( 2 \\mathclose{}\\right)\"),\n (\n \"pinv(a, (1, 0))\",\n r\"\\mathrm{pinv} \\mathopen{}\\left( a, \"\n r\"\\mathopen{}\\left( 1, 0 \\mathclose{}\\right) \\mathclose{}\\right)\",\n ),\n ],\n)\ndef test_pinv(code: str, latex: str) -> None:\n tree = ast_utils.parse_expr(code)\n assert isinstance(tree, ast.Call)\n assert expression_codegen.ExpressionCodegen().visit(tree) == latex\n\n\n# Check list for #89.\n# https://github.com/google/latexify_py/issues/89#issuecomment-1344967636\n@pytest.mark.parametrize(\n \"left,right,latex\",\n [\n (\"2\", \"3\", r\"2 \\cdot 3\"),\n (\"2\", \"y\", \"2 y\"),\n (\"2\", \"beta\", r\"2 \\beta\"),\n (\"2\", \"bar\", r\"2 \\mathrm{bar}\"),\n (\"2\", \"g(y)\", r\"2 g \\mathopen{}\\left( y \\mathclose{}\\right)\"),\n (\"2\", \"(u + v)\", r\"2 \\mathopen{}\\left( u + v \\mathclose{}\\right)\"),\n (\"x\", \"3\", r\"x \\cdot 3\"),\n (\"x\", \"y\", \"x y\"),\n (\"x\", \"beta\", r\"x \\beta\"),\n (\"x\", \"bar\", r\"x \\cdot \\mathrm{bar}\"),\n (\"x\", \"g(y)\", r\"x \\cdot g \\mathopen{}\\left( y \\mathclose{}\\right)\"),\n (\"x\", \"(u + v)\", r\"x \\cdot \\mathopen{}\\left( u + v \\mathclose{}\\right)\"),\n (\"alpha\", \"3\", r\"\\alpha \\cdot 3\"),\n (\"alpha\", \"y\", r\"\\alpha y\"),\n (\"alpha\", \"beta\", r\"\\alpha \\beta\"),\n (\"alpha\", \"bar\", r\"\\alpha \\cdot \\mathrm{bar}\"),\n (\"alpha\", \"g(y)\", r\"\\alpha \\cdot g \\mathopen{}\\left( y \\mathclose{}\\right)\"),\n (\n \"alpha\",\n \"(u + v)\",\n r\"\\alpha \\cdot \\mathopen{}\\left( u + v \\mathclose{}\\right)\",\n ),\n (\"foo\", \"3\", r\"\\mathrm{foo} \\cdot 3\"),\n (\"foo\", \"y\", r\"\\mathrm{foo} \\cdot y\"),\n (\"foo\", \"beta\", r\"\\mathrm{foo} \\cdot \\beta\"),\n (\"foo\", \"bar\", r\"\\mathrm{foo} \\cdot \\mathrm{bar}\"),\n (\n \"foo\",\n \"g(y)\",\n r\"\\mathrm{foo} \\cdot g \\mathopen{}\\left( y \\mathclose{}\\right)\",\n ),\n (\n \"foo\",\n \"(u + v)\",\n r\"\\mathrm{foo} \\cdot \\mathopen{}\\left( u + v \\mathclose{}\\right)\",\n ),\n (\"f(x)\", \"3\", r\"f \\mathopen{}\\left( x \\mathclose{}\\right) \\cdot 3\"),\n (\"f(x)\", \"y\", r\"f \\mathopen{}\\left( x \\mathclose{}\\right) \\cdot y\"),\n (\"f(x)\", \"beta\", r\"f \\mathopen{}\\left( x \\mathclose{}\\right) \\cdot \\beta\"),\n (\n \"f(x)\",\n \"bar\",\n r\"f \\mathopen{}\\left( x \\mathclose{}\\right) \\cdot \\mathrm{bar}\",\n ),\n (\n \"f(x)\",\n \"g(y)\",\n r\"f \\mathopen{}\\left( x \\mathclose{}\\right)\"\n r\" \\cdot g \\mathopen{}\\left( y \\mathclose{}\\right)\",\n ),\n (\n \"f(x)\",\n \"(u + v)\",\n r\"f \\mathopen{}\\left( x \\mathclose{}\\right)\"\n r\" \\cdot \\mathopen{}\\left( u + v \\mathclose{}\\right)\",\n ),\n (\"(s + t)\", \"3\", r\"\\mathopen{}\\left( s + t \\mathclose{}\\right) \\cdot 3\"),\n (\"(s + t)\", \"y\", r\"\\mathopen{}\\left( s + t \\mathclose{}\\right) y\"),\n (\"(s + t)\", \"beta\", r\"\\mathopen{}\\left( s + t \\mathclose{}\\right) \\beta\"),\n (\n \"(s + t)\",\n \"bar\",\n r\"\\mathopen{}\\left( s + t \\mathclose{}\\right) \\mathrm{bar}\",\n ),\n (\n \"(s + t)\",\n \"g(y)\",\n r\"\\mathopen{}\\left( s + t \\mathclose{}\\right)\"\n r\" g \\mathopen{}\\left( y \\mathclose{}\\right)\",\n ),\n (\n \"(s + t)\",\n \"(u + v)\",\n r\"\\mathopen{}\\left( s + t \\mathclose{}\\right)\"\n r\" \\mathopen{}\\left( u + v \\mathclose{}\\right)\",\n ),\n ],\n)\ndef test_remove_multiply(left: str, right: str, latex: str) -> None:\n for op in [\"*\", \"@\"]:\n tree = ast_utils.parse_expr(f\"{left} {op} {right}\")\n assert isinstance(tree, ast.BinOp)\n assert (\n expression_codegen.ExpressionCodegen(use_math_symbols=True).visit(tree)\n == latex\n )\n"
},
{
"path": "src/latexify/codegen/expression_rules.py",
"content": "\"\"\"Codegen rules for single expressions.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\nimport dataclasses\n\n# Precedences of operators for BoolOp, BinOp, UnaryOp, and Compare nodes.\n# Note that this value affects only the appearance of surrounding parentheses for each\n# expression, and does not affect the AST itself.\n# See also:\n# https://docs.python.org/3/reference/expressions.html#operator-precedence\n_PRECEDENCES: dict[type[ast.AST], int] = {\n ast.Pow: 120,\n ast.UAdd: 110,\n ast.USub: 110,\n ast.Invert: 110,\n ast.Mult: 100,\n ast.MatMult: 100,\n ast.Div: 100,\n ast.FloorDiv: 100,\n ast.Mod: 100,\n ast.Add: 90,\n ast.Sub: 90,\n ast.LShift: 80,\n ast.RShift: 80,\n ast.BitAnd: 70,\n ast.BitXor: 60,\n ast.BitOr: 50,\n ast.In: 40,\n ast.NotIn: 40,\n ast.Is: 40,\n ast.IsNot: 40,\n ast.Lt: 40,\n ast.LtE: 40,\n ast.Gt: 40,\n ast.GtE: 40,\n ast.NotEq: 40,\n ast.Eq: 40,\n # NOTE(odashi):\n # We assume that the `not` operator has the same precedence with other unary\n # operators `+`, `-` and `~`, because the LaTeX counterpart $\\lnot$ looks to have a\n # high precedence.\n # ast.Not: 30,\n ast.Not: 110,\n ast.And: 20,\n ast.Or: 10,\n}\n\n# NOTE(odashi):\n# Function invocation is treated as a unary operator with a higher precedence.\n# This ensures that the argument with a unary operator is wrapped:\n# exp(x) --> \\exp x\n# exp(-x) --> \\exp (-x)\n# -exp(x) --> - \\exp x\n_CALL_PRECEDENCE = _PRECEDENCES[ast.UAdd] + 1\n\n_INF_PRECEDENCE = 1_000_000\n\n\ndef get_precedence(node: ast.AST) -> int:\n \"\"\"Obtains the precedence of the subtree.\n\n Args:\n node: Subtree to investigate.\n\n Returns:\n If `node` is a subtree with some operator, returns the precedence of the\n operator. Otherwise, returns a number larger enough from other precedences.\n \"\"\"\n if isinstance(node, ast.Call):\n return _CALL_PRECEDENCE\n\n if isinstance(node, (ast.BinOp, ast.UnaryOp, ast.BoolOp)):\n return _PRECEDENCES[type(node.op)]\n\n if isinstance(node, ast.Compare):\n # Compare operators have the same precedence. It is enough to check only the\n # first operator.\n return _PRECEDENCES[type(node.ops[0])]\n\n return _INF_PRECEDENCE\n\n\n@dataclasses.dataclass(frozen=True)\nclass BinOperandRule:\n \"\"\"Syntax rules for operands of BinOp.\"\"\"\n\n # Whether to require wrapping operands by parentheses according to the precedence.\n wrap: bool = True\n\n # Whether to require wrapping operands by parentheses if the operand has the same\n # precedence with this operator.\n # This is used to control the behavior of non-associative operators.\n force: bool = False\n\n\n@dataclasses.dataclass(frozen=True)\nclass BinOpRule:\n \"\"\"Syntax rules for BinOp.\"\"\"\n\n # Left/middle/right syntaxes to wrap operands.\n latex_left: str\n latex_middle: str\n latex_right: str\n\n # Operand rules.\n operand_left: BinOperandRule = dataclasses.field(default_factory=BinOperandRule)\n operand_right: BinOperandRule = dataclasses.field(default_factory=BinOperandRule)\n\n # Whether to assume the resulting syntax is wrapped by some bracket operators.\n # If True, the parent operator can avoid wrapping this operator by parentheses.\n is_wrapped: bool = False\n\n\nBIN_OP_RULES: dict[type[ast.operator], BinOpRule] = {\n ast.Pow: BinOpRule(\n \"\",\n \"^{\",\n \"}\",\n operand_left=BinOperandRule(force=True),\n operand_right=BinOperandRule(wrap=False),\n ),\n ast.Mult: BinOpRule(\"\", r\" \\cdot \", \"\"),\n ast.MatMult: BinOpRule(\"\", r\" \\cdot \", \"\"),\n ast.Div: BinOpRule(\n r\"\\frac{\",\n \"}{\",\n \"}\",\n operand_left=BinOperandRule(wrap=False),\n operand_right=BinOperandRule(wrap=False),\n ),\n ast.FloorDiv: BinOpRule(\n r\"\\left\\lfloor\\frac{\",\n \"}{\",\n r\"}\\right\\rfloor\",\n operand_left=BinOperandRule(wrap=False),\n operand_right=BinOperandRule(wrap=False),\n is_wrapped=True,\n ),\n ast.Mod: BinOpRule(\n \"\", r\" \\mathbin{\\%} \", \"\", operand_right=BinOperandRule(force=True)\n ),\n ast.Add: BinOpRule(\"\", \" + \", \"\"),\n ast.Sub: BinOpRule(\"\", \" - \", \"\", operand_right=BinOperandRule(force=True)),\n ast.LShift: BinOpRule(\"\", r\" \\ll \", \"\", operand_right=BinOperandRule(force=True)),\n ast.RShift: BinOpRule(\"\", r\" \\gg \", \"\", operand_right=BinOperandRule(force=True)),\n ast.BitAnd: BinOpRule(\"\", r\" \\mathbin{\\&} \", \"\"),\n ast.BitXor: BinOpRule(\"\", r\" \\oplus \", \"\"),\n ast.BitOr: BinOpRule(\"\", r\" \\mathbin{|} \", \"\"),\n}\n\n# Typeset for BinOp of sets.\nSET_BIN_OP_RULES: dict[type[ast.operator], BinOpRule] = {\n **BIN_OP_RULES,\n ast.Sub: BinOpRule(\n \"\", r\" \\setminus \", \"\", operand_right=BinOperandRule(force=True)\n ),\n ast.BitAnd: BinOpRule(\"\", r\" \\cap \", \"\"),\n ast.BitXor: BinOpRule(\"\", r\" \\mathbin{\\triangle} \", \"\"),\n ast.BitOr: BinOpRule(\"\", r\" \\cup \", \"\"),\n}\n\nUNARY_OPS: dict[type[ast.unaryop], str] = {\n ast.Invert: r\"\\mathord{\\sim} \",\n ast.UAdd: \"+\", # Explicitly adds the $+$ operator.\n ast.USub: \"-\",\n ast.Not: r\"\\lnot \",\n}\n\nCOMPARE_OPS: dict[type[ast.cmpop], str] = {\n ast.Eq: \"=\",\n ast.Gt: \">\",\n ast.GtE: r\"\\ge\",\n ast.In: r\"\\in\",\n ast.Is: r\"\\equiv\",\n ast.IsNot: r\"\\not\\equiv\",\n ast.Lt: \"<\",\n ast.LtE: r\"\\le\",\n ast.NotEq: r\"\\ne\",\n ast.NotIn: r\"\\notin\",\n}\n\n# Typeset for Compare of sets.\nSET_COMPARE_OPS: dict[type[ast.cmpop], str] = {\n **COMPARE_OPS,\n ast.Gt: r\"\\supset\",\n ast.GtE: r\"\\supseteq\",\n ast.Lt: r\"\\subset\",\n ast.LtE: r\"\\subseteq\",\n}\n\nBOOL_OPS: dict[type[ast.boolop], str] = {\n ast.And: r\"\\land\",\n ast.Or: r\"\\lor\",\n}\n\n\n@dataclasses.dataclass(frozen=True)\nclass FunctionRule:\n \"\"\"Codegen rules for functions.\n\n Attributes:\n left: LaTeX expression concatenated to the left-hand side of the arguments.\n right: LaTeX expression concatenated to the right-hand side of the arguments.\n is_unary: Whether the function is treated as a unary operator or not.\n is_wrapped: Whether the resulting syntax is wrapped by brackets or not.\n \"\"\"\n\n left: str\n right: str = \"\"\n is_unary: bool = False\n is_wrapped: bool = False\n\n\n# name => left_syntax, right_syntax, is_wrapped\nBUILTIN_FUNCS: dict[str, FunctionRule] = {\n \"abs\": FunctionRule(r\"\\mathopen{}\\left|\", r\"\\mathclose{}\\right|\", is_wrapped=True),\n \"acos\": FunctionRule(r\"\\arccos\", is_unary=True),\n \"acosh\": FunctionRule(r\"\\mathrm{arcosh}\", is_unary=True),\n \"arccos\": FunctionRule(r\"\\arccos\", is_unary=True),\n \"arccot\": FunctionRule(r\"\\mathrm{arccot}\", is_unary=True),\n \"arccsc\": FunctionRule(r\"\\mathrm{arccsc}\", is_unary=True),\n \"arcosh\": FunctionRule(r\"\\mathrm{arcosh}\", is_unary=True),\n \"arcoth\": FunctionRule(r\"\\mathrm{arcoth}\", is_unary=True),\n \"arcsec\": FunctionRule(r\"\\mathrm{arcsec}\", is_unary=True),\n \"arcsch\": FunctionRule(r\"\\mathrm{arcsch}\", is_unary=True),\n \"arcsin\": FunctionRule(r\"\\arcsin\", is_unary=True),\n \"arctan\": FunctionRule(r\"\\arctan\", is_unary=True),\n \"arsech\": FunctionRule(r\"\\mathrm{arsech}\", is_unary=True),\n \"arsinh\": FunctionRule(r\"\\mathrm{arsinh}\", is_unary=True),\n \"artanh\": FunctionRule(r\"\\mathrm{artanh}\", is_unary=True),\n \"asin\": FunctionRule(r\"\\arcsin\", is_unary=True),\n \"asinh\": FunctionRule(r\"\\mathrm{arsinh}\", is_unary=True),\n \"atan\": FunctionRule(r\"\\arctan\", is_unary=True),\n \"atanh\": FunctionRule(r\"\\mathrm{artanh}\", is_unary=True),\n \"ceil\": FunctionRule(\n r\"\\mathopen{}\\left\\lceil\", r\"\\mathclose{}\\right\\rceil\", is_wrapped=True\n ),\n \"cos\": FunctionRule(r\"\\cos\", is_unary=True),\n \"cosh\": FunctionRule(r\"\\cosh\", is_unary=True),\n \"cot\": FunctionRule(r\"\\cot\", is_unary=True),\n \"coth\": FunctionRule(r\"\\coth\", is_unary=True),\n \"csc\": FunctionRule(r\"\\csc\", is_unary=True),\n \"csch\": FunctionRule(r\"\\mathrm{csch}\", is_unary=True),\n \"exp\": FunctionRule(r\"\\exp\", is_unary=True),\n \"fabs\": FunctionRule(r\"\\mathopen{}\\left|\", r\"\\mathclose{}\\right|\", is_wrapped=True),\n \"factorial\": FunctionRule(\"\", \"!\", is_unary=True),\n \"floor\": FunctionRule(\n r\"\\mathopen{}\\left\\lfloor\", r\"\\mathclose{}\\right\\rfloor\", is_wrapped=True\n ),\n \"fsum\": FunctionRule(r\"\\sum\", is_unary=True),\n \"gamma\": FunctionRule(r\"\\Gamma\"),\n \"log\": FunctionRule(r\"\\log\", is_unary=True),\n \"log10\": FunctionRule(r\"\\log_{10}\", is_unary=True),\n \"log2\": FunctionRule(r\"\\log_2\", is_unary=True),\n \"prod\": FunctionRule(r\"\\prod\", is_unary=True),\n \"sec\": FunctionRule(r\"\\sec\", is_unary=True),\n \"sech\": FunctionRule(r\"\\mathrm{sech}\", is_unary=True),\n \"sin\": FunctionRule(r\"\\sin\", is_unary=True),\n \"sinh\": FunctionRule(r\"\\sinh\", is_unary=True),\n \"sqrt\": FunctionRule(r\"\\sqrt{\", \"}\", is_wrapped=True),\n \"sum\": FunctionRule(r\"\\sum\", is_unary=True),\n \"tan\": FunctionRule(r\"\\tan\", is_unary=True),\n \"tanh\": FunctionRule(r\"\\tanh\", is_unary=True),\n}\n\nMATH_SYMBOLS = {\n \"aleph\",\n \"alpha\",\n \"beta\",\n \"beth\",\n \"chi\",\n \"daleth\",\n \"delta\",\n \"digamma\",\n \"epsilon\",\n \"eta\",\n \"gamma\",\n \"gimel\",\n \"hbar\",\n \"infty\",\n \"iota\",\n \"kappa\",\n \"lambda\",\n \"mu\",\n \"nabla\",\n \"nu\",\n \"omega\",\n \"phi\",\n \"pi\",\n \"psi\",\n \"rho\",\n \"sigma\",\n \"tau\",\n \"theta\",\n \"upsilon\",\n \"varepsilon\",\n \"varkappa\",\n \"varphi\",\n \"varpi\",\n \"varrho\",\n \"varsigma\",\n \"vartheta\",\n \"xi\",\n \"zeta\",\n \"Delta\",\n \"Gamma\",\n \"Lambda\",\n \"Omega\",\n \"Phi\",\n \"Pi\",\n \"Psi\",\n \"Sigma\",\n \"Theta\",\n \"Upsilon\",\n \"Xi\",\n}\n"
},
{
"path": "src/latexify/codegen/expression_rules_test.py",
"content": "\"\"\"Tests for latexify.codegen.expression_rules.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\n\nimport pytest\n\nfrom latexify.codegen import expression_rules\n\n\n@pytest.mark.parametrize(\n \"node,precedence\",\n [\n (\n ast.Call(func=ast.Name(id=\"func\", ctx=ast.Load()), args=[], keywords=[]),\n expression_rules._CALL_PRECEDENCE,\n ),\n (\n ast.BinOp(\n left=ast.Name(id=\"left\", ctx=ast.Load()),\n op=ast.Add(),\n right=ast.Name(id=\"right\", ctx=ast.Load()),\n ),\n expression_rules._PRECEDENCES[ast.Add],\n ),\n (\n ast.UnaryOp(op=ast.UAdd(), operand=ast.Name(id=\"operand\", ctx=ast.Load())),\n expression_rules._PRECEDENCES[ast.UAdd],\n ),\n (\n ast.BoolOp(op=ast.And(), values=[ast.Name(id=\"value\", ctx=ast.Load())]),\n expression_rules._PRECEDENCES[ast.And],\n ),\n (\n ast.Compare(\n left=ast.Name(id=\"left\", ctx=ast.Load()),\n ops=[ast.Eq()],\n comparators=[ast.Name(id=\"right\", ctx=ast.Load())],\n ),\n expression_rules._PRECEDENCES[ast.Eq],\n ),\n (ast.Name(id=\"name\", ctx=ast.Load()), expression_rules._INF_PRECEDENCE),\n (\n ast.Attribute(\n value=ast.Name(id=\"value\", ctx=ast.Load()), attr=\"attr\", ctx=ast.Load()\n ),\n expression_rules._INF_PRECEDENCE,\n ),\n ],\n)\ndef test_get_precedence(node: ast.AST, precedence: int) -> None:\n assert expression_rules.get_precedence(node) == precedence\n"
},
{
"path": "src/latexify/codegen/function_codegen.py",
"content": "\"\"\"Codegen for single functions.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\nimport sys\n\nfrom latexify import ast_utils, exceptions\nfrom latexify.codegen import codegen_utils, expression_codegen, identifier_converter\n\n\nclass FunctionCodegen(ast.NodeVisitor):\n \"\"\"Codegen for single functions.\n\n This codegen works for Module with single FunctionDef node to generate a single\n LaTeX expression of the given function.\n \"\"\"\n\n _identifier_converter: identifier_converter.IdentifierConverter\n _use_signature: bool\n\n def __init__(\n self,\n *,\n use_math_symbols: bool = False,\n use_signature: bool = True,\n use_set_symbols: bool = False,\n escape_underscores: bool = True,\n ) -> None:\n \"\"\"Initializer.\n\n Args:\n use_math_symbols: Whether to convert identifiers with a math symbol surface\n (e.g., \"alpha\") to the LaTeX symbol (e.g., \"\\\\alpha\").\n use_signature: Whether to add the function signature before the expression\n or not.\n use_set_symbols: Whether to use set symbols or not.\n \"\"\"\n self._expression_codegen = expression_codegen.ExpressionCodegen(\n use_math_symbols=use_math_symbols,\n use_set_symbols=use_set_symbols,\n escape_underscores=escape_underscores,\n )\n self._identifier_converter = identifier_converter.IdentifierConverter(\n use_math_symbols=use_math_symbols, escape_underscores=escape_underscores\n )\n self._use_signature = use_signature\n\n def generic_visit(self, node: ast.AST) -> str:\n raise exceptions.LatexifyNotSupportedError(\n f\"Unsupported AST: {type(node).__name__}\"\n )\n\n def visit_Module(self, node: ast.Module) -> str:\n \"\"\"Visit a Module node.\"\"\"\n return self.visit(node.body[0])\n\n def visit_FunctionDef(self, node: ast.FunctionDef) -> str:\n \"\"\"Visit a FunctionDef node.\"\"\"\n # Function name\n name_str = self._identifier_converter.convert(node.name)[0]\n\n # Arguments\n arg_strs = [\n self._identifier_converter.convert(arg.arg)[0] for arg in node.args.args\n ]\n\n body_strs: list[str] = []\n\n # Assignment statements (if any): x = ...\n for child in node.body[:-1]:\n if isinstance(child, ast.Expr) and ast_utils.is_constant(child.value):\n continue\n\n if not isinstance(child, ast.Assign):\n raise exceptions.LatexifyNotSupportedError(\n \"Codegen supports only Assign nodes in multiline functions, \"\n f\"but got: {type(child).__name__}\"\n )\n body_strs.append(self.visit(child))\n\n return_stmt = node.body[-1]\n\n if sys.version_info.minor >= 10:\n if not isinstance(return_stmt, (ast.Return, ast.If, ast.Match)):\n raise exceptions.LatexifySyntaxError(\n f\"Unsupported last statement: {type(return_stmt).__name__}\"\n )\n else:\n if not isinstance(return_stmt, (ast.Return, ast.If)):\n raise exceptions.LatexifySyntaxError(\n f\"Unsupported last statement: {type(return_stmt).__name__}\"\n )\n\n # Function signature: f(x, ...)\n signature_str = name_str + \"(\" + \", \".join(arg_strs) + \")\"\n\n # Function definition: f(x, ...) \\triangleq ...\n return_str = self.visit(return_stmt)\n if self._use_signature:\n return_str = signature_str + \" = \" + return_str\n\n if not body_strs:\n # Only the definition.\n return return_str\n\n # Definition with several assignments. Wrap all statements with array.\n body_strs.append(return_str)\n return r\"\\begin{array}{l} \" + r\" \\\\ \".join(body_strs) + r\" \\end{array}\"\n\n def visit_Assign(self, node: ast.Assign) -> str:\n \"\"\"Visit an Assign node.\"\"\"\n operands: list[str] = [self._expression_codegen.visit(t) for t in node.targets]\n operands.append(self._expression_codegen.visit(node.value))\n return \" = \".join(operands)\n\n def visit_Return(self, node: ast.Return) -> str:\n \"\"\"Visit a Return node.\"\"\"\n return (\n self._expression_codegen.visit(node.value)\n if node.value is not None\n else codegen_utils.convert_constant(None)\n )\n\n def visit_If(self, node: ast.If) -> str:\n \"\"\"Visit an If node.\"\"\"\n latex = r\"\\left\\{ \\begin{array}{ll} \"\n\n current_stmt: ast.stmt = node\n\n while isinstance(current_stmt, ast.If):\n if len(current_stmt.body) != 1 or len(current_stmt.orelse) != 1:\n raise exceptions.LatexifySyntaxError(\n \"Multiple statements are not supported in If nodes.\"\n )\n\n cond_latex = self._expression_codegen.visit(current_stmt.test)\n true_latex = self.visit(current_stmt.body[0])\n latex += true_latex + r\", & \\mathrm{if} \\ \" + cond_latex + r\" \\\\ \"\n current_stmt = current_stmt.orelse[0]\n\n latex += self.visit(current_stmt)\n return latex + r\", & \\mathrm{otherwise} \\end{array} \\right.\"\n\n def visit_Match(self, node: ast.Match) -> str:\n \"\"\"Visit a Match node\"\"\"\n if not (\n len(node.cases) >= 2\n and isinstance(node.cases[-1].pattern, ast.MatchAs)\n and node.cases[-1].pattern.name is None\n ):\n raise exceptions.LatexifySyntaxError(\n \"Match statement must contain the wildcard.\"\n )\n\n subject_latex = self._expression_codegen.visit(node.subject)\n case_latexes: list[str] = []\n\n for i, case in enumerate(node.cases):\n if len(case.body) != 1 or not isinstance(case.body[0], ast.Return):\n raise exceptions.LatexifyNotSupportedError(\n \"Match cases must contain exactly 1 return statement.\"\n )\n\n if i < len(node.cases) - 1:\n body_latex = self.visit(case.body[0])\n cond_latex = self.visit(case.pattern)\n case_latexes.append(\n body_latex + r\", & \\mathrm{if} \\ \" + subject_latex + cond_latex\n )\n else:\n case_latexes.append(\n self.visit(node.cases[-1].body[0]) + r\", & \\mathrm{otherwise}\"\n )\n\n return (\n r\"\\left\\{ \\begin{array}{ll} \"\n + r\" \\\\ \".join(case_latexes)\n + r\" \\end{array} \\right.\"\n )\n\n def visit_MatchValue(self, node: ast.MatchValue) -> str:\n \"\"\"Visit a MatchValue node\"\"\"\n latex = self._expression_codegen.visit(node.value)\n return \" = \" + latex\n"
},
{
"path": "src/latexify/codegen/function_codegen_match_test.py",
"content": "\"\"\"Tests for FunctionCodegen with match statements.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\nimport textwrap\n\nimport pytest\n\nfrom latexify import exceptions, test_utils\nfrom latexify.codegen import function_codegen\n\n\n@test_utils.require_at_least(10)\ndef test_functiondef_match() -> None:\n tree = ast.parse(\n textwrap.dedent(\n \"\"\"\n def f(x):\n match x:\n case 0:\n return 1\n case _:\n return 3 * x\n \"\"\"\n )\n )\n expected = (\n r\"f(x) =\"\n r\" \\left\\{ \\begin{array}{ll}\"\n r\" 1, & \\mathrm{if} \\ x = 0 \\\\\"\n r\" 3 x, & \\mathrm{otherwise}\"\n r\" \\end{array} \\right.\"\n )\n assert function_codegen.FunctionCodegen().visit(tree) == expected\n\n\n@test_utils.require_at_least(10)\ndef test_matchvalue() -> None:\n tree = ast.parse(\n textwrap.dedent(\n \"\"\"\n match x:\n case 0:\n return 1\n case _:\n return 2\n \"\"\"\n )\n ).body[0]\n expected = (\n r\"\\left\\{ \\begin{array}{ll}\"\n r\" 1, & \\mathrm{if} \\ x = 0 \\\\\"\n r\" 2, & \\mathrm{otherwise}\"\n r\" \\end{array} \\right.\"\n )\n assert function_codegen.FunctionCodegen().visit(tree) == expected\n\n\n@test_utils.require_at_least(10)\ndef test_multiple_matchvalue() -> None:\n tree = ast.parse(\n textwrap.dedent(\n \"\"\"\n match x:\n case 0:\n return 1\n case 1:\n return 2\n case _:\n return 3\n \"\"\"\n )\n ).body[0]\n expected = (\n r\"\\left\\{ \\begin{array}{ll}\"\n r\" 1, & \\mathrm{if} \\ x = 0 \\\\\"\n r\" 2, & \\mathrm{if} \\ x = 1 \\\\\"\n r\" 3, & \\mathrm{otherwise}\"\n r\" \\end{array} \\right.\"\n )\n assert function_codegen.FunctionCodegen().visit(tree) == expected\n\n\n@test_utils.require_at_least(10)\ndef test_single_matchvalue_no_wildcards() -> None:\n tree = ast.parse(\n textwrap.dedent(\n \"\"\"\n match x:\n case 0:\n return 1\n \"\"\"\n )\n ).body[0]\n\n with pytest.raises(\n exceptions.LatexifySyntaxError,\n match=r\"^Match statement must contain the wildcard\\.$\",\n ):\n function_codegen.FunctionCodegen().visit(tree)\n\n\n@test_utils.require_at_least(10)\ndef test_multiple_matchvalue_no_wildcards() -> None:\n tree = ast.parse(\n textwrap.dedent(\n \"\"\"\n match x:\n case 0:\n return 1\n case 1:\n return 2\n \"\"\"\n )\n ).body[0]\n\n with pytest.raises(\n exceptions.LatexifySyntaxError,\n match=r\"^Match statement must contain the wildcard\\.$\",\n ):\n function_codegen.FunctionCodegen().visit(tree)\n\n\n@test_utils.require_at_least(10)\ndef test_matchas_nonempty() -> None:\n tree = ast.parse(\n textwrap.dedent(\n \"\"\"\n match x:\n case [x] as y:\n return 1\n case _:\n return 2\n \"\"\"\n )\n ).body[0]\n\n with pytest.raises(\n exceptions.LatexifyNotSupportedError,\n match=r\"^Unsupported AST: MatchAs$\",\n ):\n function_codegen.FunctionCodegen().visit(tree)\n\n\n@test_utils.require_at_least(10)\ndef test_matchvalue_no_return() -> None:\n tree = ast.parse(\n textwrap.dedent(\n \"\"\"\n match x:\n case 0:\n x = 5\n case _:\n return 0\n \"\"\"\n )\n ).body[0]\n\n with pytest.raises(\n exceptions.LatexifyNotSupportedError,\n match=r\"^Match cases must contain exactly 1 return statement\\.$\",\n ):\n function_codegen.FunctionCodegen().visit(tree)\n\n\n@test_utils.require_at_least(10)\ndef test_matchvalue_mutliple_statements() -> None:\n tree = ast.parse(\n textwrap.dedent(\n \"\"\"\n match x:\n case 0:\n x = 5\n return 1\n case _:\n return 0\n \"\"\"\n )\n ).body[0]\n\n with pytest.raises(\n exceptions.LatexifyNotSupportedError,\n match=r\"^Match cases must contain exactly 1 return statement\\.$\",\n ):\n function_codegen.FunctionCodegen().visit(tree)\n"
},
{
"path": "src/latexify/codegen/function_codegen_test.py",
"content": "\"\"\"Tests for latexify.codegen.function_codegen.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\nimport textwrap\n\nimport pytest\n\nfrom latexify import exceptions\nfrom latexify.codegen import function_codegen\n\n\ndef test_generic_visit() -> None:\n class UnknownNode(ast.AST):\n pass\n\n with pytest.raises(\n exceptions.LatexifyNotSupportedError,\n match=r\"^Unsupported AST: UnknownNode$\",\n ):\n function_codegen.FunctionCodegen().visit(UnknownNode())\n\n\ndef test_visit_functiondef_use_signature() -> None:\n tree = ast.parse(\n textwrap.dedent(\n \"\"\"\n def f(x):\n return x\n \"\"\"\n )\n ).body[0]\n assert isinstance(tree, ast.FunctionDef)\n\n latex_without_flag = \"x\"\n latex_with_flag = r\"f(x) = x\"\n assert function_codegen.FunctionCodegen().visit(tree) == latex_with_flag\n assert (\n function_codegen.FunctionCodegen(use_signature=False).visit(tree)\n == latex_without_flag\n )\n assert (\n function_codegen.FunctionCodegen(use_signature=True).visit(tree)\n == latex_with_flag\n )\n\n\ndef test_visit_functiondef_ignore_docstring() -> None:\n tree = ast.parse(\n textwrap.dedent(\n \"\"\"\n def f(x):\n '''docstring'''\n return x\n \"\"\"\n )\n ).body[0]\n assert isinstance(tree, ast.FunctionDef)\n\n latex = r\"f(x) = x\"\n assert function_codegen.FunctionCodegen().visit(tree) == latex\n\n\ndef test_visit_functiondef_ignore_multiple_constants() -> None:\n tree = ast.parse(\n textwrap.dedent(\n \"\"\"\n def f(x):\n '''docstring'''\n 3\n True\n return x\n \"\"\"\n )\n ).body[0]\n assert isinstance(tree, ast.FunctionDef)\n\n latex = r\"f(x) = x\"\n assert function_codegen.FunctionCodegen().visit(tree) == latex\n"
},
{
"path": "src/latexify/codegen/identifier_converter.py",
"content": "\"\"\"Utility to convert identifiers.\"\"\"\n\nfrom __future__ import annotations\n\nfrom latexify.codegen import expression_rules\n\n\nclass IdentifierConverter:\n r\"\"\"Converts Python identifiers to appropriate LaTeX expression.\n\n This converter applies following rules:\n - `foo` --> `\\foo`, if `use_math_symbols == True` and the given identifier\n matches a supported math symbol name.\n - `x` --> `x`, if the given identifier is exactly 1 character (except `_`)\n - `foo_bar` --> `\\mathrm{foo\\_bar}`, otherwise.\n \"\"\"\n\n _use_math_symbols: bool\n _use_mathrm: bool\n _escape_underscores: bool\n\n def __init__(\n self,\n *,\n use_math_symbols: bool,\n use_mathrm: bool = True,\n escape_underscores: bool = True,\n ) -> None:\n r\"\"\"Initializer.\n\n Args:\n use_math_symbols: Whether to convert identifiers with math symbol names to\n appropriate LaTeX command.\n use_mathrm: Whether to wrap the resulting expression by \\mathrm, if\n applicable.\n escape_underscores: Whether to prefix any underscores in identifiers with\n '\\\\', disable to allow subscripts in generated latex.\n \"\"\"\n self._use_math_symbols = use_math_symbols\n self._use_mathrm = use_mathrm\n self._escape_underscores = escape_underscores\n\n def convert(self, name: str) -> tuple[str, bool]:\n \"\"\"Converts Python identifier to LaTeX expression.\n\n Args:\n name: Identifier name.\n\n Returns:\n Tuple of following values:\n - latex: Corresponding LaTeX expression.\n - is_single_character: Whether `latex` can be treated as a single\n character or not.\n Raises:\n LatexifyError: Resulting latex is not valid. This most likely occurs where\n the symbol starts or ends with an underscore, and escape_underscores=False.\n \"\"\"\n if not self._escape_underscores and \"_\" in name:\n # Check if we are going to generate an invalid Latex string. Better to\n # raise an exception here than have the resulting Latex fail to\n # compile/display\n name_splits = name.split(\"_\")\n if not all(name_splits):\n raise ValueError(\n \"Neither preceding/trailing underscores nor double underscores is \"\n f\"allowed by the `escape_underscores` option, but got: {name}\"\n )\n elems = [\n IdentifierConverter(\n use_math_symbols=self._use_math_symbols,\n use_mathrm=False,\n escape_underscores=True,\n ).convert(n)[0]\n for n in name_splits\n ]\n # Wrap sub identifiers in nested braces\n name = \"_{\".join(elems) + \"}\" * (len(elems) - 1)\n\n if self._use_math_symbols and name in expression_rules.MATH_SYMBOLS:\n return \"\\\\\" + name, True\n\n if len(name) == 1 and name != \"_\":\n return name, True\n\n escaped = name.replace(\"_\", r\"\\_\") if self._escape_underscores else name\n wrapped = rf\"\\mathrm{{{escaped}}}\" if self._use_mathrm else escaped\n\n return wrapped, False\n"
},
{
"path": "src/latexify/codegen/identifier_converter_test.py",
"content": "\"\"\"Tests for latexify.codegen.identifier_converter.\"\"\"\n\nfrom __future__ import annotations\n\nimport pytest\n\nfrom latexify.codegen import identifier_converter\n\n\n@pytest.mark.parametrize(\n \"name,use_math_symbols,use_mathrm,escape_underscores,expected\",\n [\n (\"a\", False, True, True, (\"a\", True)),\n (\"_\", False, True, True, (r\"\\mathrm{\\_}\", False)),\n (\"aa\", False, True, True, (r\"\\mathrm{aa}\", False)),\n (\"a1\", False, True, True, (r\"\\mathrm{a1}\", False)),\n (\"a_\", False, True, True, (r\"\\mathrm{a\\_}\", False)),\n (\"_a\", False, True, True, (r\"\\mathrm{\\_a}\", False)),\n (\"_1\", False, True, True, (r\"\\mathrm{\\_1}\", False)),\n (\"__\", False, True, True, (r\"\\mathrm{\\_\\_}\", False)),\n (\"a_a\", False, True, True, (r\"\\mathrm{a\\_a}\", False)),\n (\"a__\", False, True, True, (r\"\\mathrm{a\\_\\_}\", False)),\n (\"a_1\", False, True, True, (r\"\\mathrm{a\\_1}\", False)),\n (\"alpha\", False, True, True, (r\"\\mathrm{alpha}\", False)),\n (\"alpha\", True, True, True, (r\"\\alpha\", True)),\n (\"alphabet\", True, True, True, (r\"\\mathrm{alphabet}\", False)),\n (\"foo\", False, True, True, (r\"\\mathrm{foo}\", False)),\n (\"foo\", True, True, True, (r\"\\mathrm{foo}\", False)),\n (\"foo\", True, False, True, (r\"foo\", False)),\n (\"aa\", False, True, False, (r\"\\mathrm{aa}\", False)),\n (\"a_a\", False, True, False, (r\"\\mathrm{a_{a}}\", False)),\n (\"a_1\", False, True, False, (r\"\\mathrm{a_{1}}\", False)),\n (\"alpha\", True, False, False, (r\"\\alpha\", True)),\n (\"alpha_1\", True, False, False, (r\"\\alpha_{1}\", False)),\n (\"x_alpha\", True, False, False, (r\"x_{\\alpha}\", False)),\n (\"x_alpha_beta\", True, False, False, (r\"x_{\\alpha_{\\beta}}\", False)),\n (\"alpha_beta\", True, False, False, (r\"\\alpha_{\\beta}\", False)),\n ],\n)\ndef test_identifier_converter(\n name: str,\n use_math_symbols: bool,\n use_mathrm: bool,\n escape_underscores: bool,\n expected: tuple[str, bool],\n) -> None:\n assert (\n identifier_converter.IdentifierConverter(\n use_math_symbols=use_math_symbols,\n use_mathrm=use_mathrm,\n escape_underscores=escape_underscores,\n ).convert(name)\n == expected\n )\n\n\n@pytest.mark.parametrize(\n \"name,use_math_symbols,use_mathrm,escape_underscores\",\n [\n (\"_\", False, True, False),\n (\"a_\", False, True, False),\n (\"_a\", False, True, False),\n (\"_1\", False, True, False),\n (\"__\", False, True, False),\n (\"a__\", False, True, False),\n (\"alpha_\", True, False, False),\n (\"_alpha\", True, False, False),\n (\"x__alpha\", True, False, False),\n ],\n)\ndef test_identifier_converter_failure(\n name: str,\n use_math_symbols: bool,\n use_mathrm: bool,\n escape_underscores: bool,\n) -> None:\n with pytest.raises(ValueError):\n identifier_converter.IdentifierConverter(\n use_math_symbols=use_math_symbols,\n use_mathrm=use_mathrm,\n escape_underscores=escape_underscores,\n ).convert(name)\n"
},
{
"path": "src/latexify/codegen/latex.py",
"content": "\"\"\"Definition of Latex.\"\"\"\n\nfrom __future__ import annotations\n\nfrom collections.abc import Iterable\nfrom typing import Union\n\nLatexLike = Union[str, \"Latex\"]\n\n\nclass Latex:\n \"\"\"LaTeX expression string for ease of writing the codegen source.\"\"\"\n\n _raw: str\n\n def __init__(self, raw: str) -> None:\n \"\"\"Initializer.\n\n Args:\n raw: Direct string of the underlying expression.\n \"\"\"\n self._raw = raw\n\n def __eq__(self, other: object) -> bool:\n \"\"\"Checks equality.\n\n Args:\n other: Other object to check equality.\n\n Returns:\n True if other is Latex and the underlying expression is the same as self,\n False otherwise.\n \"\"\"\n return isinstance(other, Latex) and other._raw == self._raw\n\n def __str__(self) -> str:\n \"\"\"Returns the underlying expression.\n\n Returns:\n The underlying expression.\n \"\"\"\n return self._raw\n\n def __add__(self, other: object) -> Latex:\n \"\"\"Concatenates two expressions.\n\n Args:\n other: The expression to be concatenated to the right side of self.\n\n Returns:\n A new expression: \"{self}{other}\"\n \"\"\"\n if isinstance(other, str):\n return Latex(self._raw + other)\n if isinstance(other, Latex):\n return Latex(self._raw + other._raw)\n raise ValueError(\"Unsupported operation.\")\n\n def __radd__(self, other: object) -> Latex:\n \"\"\"Concatenates two expressions.\n\n Args:\n other: The expression to be concatenated to the left side of self.\n\n Returns:\n A new expression: \"{other}{self}\"\n \"\"\"\n if isinstance(other, str):\n return Latex(other + self._raw)\n if isinstance(other, Latex):\n return Latex(other._raw + self._raw)\n raise ValueError(\"Unsupported operation.\")\n\n @staticmethod\n def opt(src: LatexLike) -> Latex:\n \"\"\"Wraps the expression by \"[\" and \"]\".\n\n This wrapping is used when the expression needs to be wrapped as an optional\n argument of the environment.\n\n Args:\n src: Original expression.\n\n Returns:\n A new expression with surrounding brackets.\n \"\"\"\n return Latex(\"[\" + str(src) + \"]\")\n\n @staticmethod\n def arg(src: LatexLike) -> Latex:\n \"\"\"Wraps the expression by \"{\" and \"}\".\n\n This wrapping is used when the expression needs to be wrapped as an argument of\n other expressions.\n\n Args:\n src: Original expression.\n\n Returns:\n A new expression with surrounding brackets.\n \"\"\"\n return Latex(\"{\" + str(src) + \"}\")\n\n @staticmethod\n def paren(src: LatexLike) -> Latex:\n \"\"\"Adds surrounding parentheses: \"(\" and \")\".\n\n Args:\n src: Original expression.\n\n Returns:\n A new expression with surrounding brackets.\n \"\"\"\n return Latex(r\"\\mathopen{}\\left( \" + str(src) + r\" \\mathclose{}\\right)\")\n\n @staticmethod\n def curly(src: LatexLike) -> Latex:\n \"\"\"Adds surrounding curly brackets: \"\\\\{\" and \"\\\\}\".\n\n Args:\n src: Original expression.\n\n Returns:\n A new expression with surrounding brackets.\n \"\"\"\n return Latex(r\"\\mathopen{}\\left\\{ \" + str(src) + r\" \\mathclose{}\\right\\}\")\n\n @staticmethod\n def square(src: LatexLike) -> Latex:\n \"\"\"Adds surrounding square brackets: \"[\" and \"]\".\n\n Args:\n src: Original expression.\n\n Returns:\n A new expression with surrounding brackets.\n \"\"\"\n return Latex(r\"\\mathopen{}\\left[ \" + str(src) + r\" \\mathclose{}\\right]\")\n\n @staticmethod\n def command(\n name: str,\n *,\n options: list[LatexLike] | None = None,\n args: list[LatexLike] | None = None,\n ) -> Latex:\n \"\"\"Makes a Latex command expression.\n\n Args:\n name: Name of the command.\n options: List of optional arguments.\n args: List of arguments.\n\n Returns:\n A new expression.\n \"\"\"\n elms: list[LatexLike] = [rf\"\\{name}\"]\n if options is not None:\n elms += [Latex.opt(x) for x in options]\n if args is not None:\n elms += [Latex.arg(x) for x in args]\n\n return Latex.join(\"\", elms)\n\n @staticmethod\n def environment(\n name: str,\n *,\n options: list[LatexLike] | None = None,\n args: list[LatexLike] | None = None,\n content: LatexLike | None = None,\n ) -> Latex:\n \"\"\"Makes a Latex environment expression.\n\n Args:\n name: Name of the environment.\n options: List of optional arguments.\n args: List of arguments.\n content: Inner content of the environment.\n\n Returns:\n A new expression.\n \"\"\"\n begin_elms: list[LatexLike] = [rf\"\\begin{{{name}}}\"]\n if options is not None:\n begin_elms += [Latex.opt(x) for x in options]\n if args is not None:\n begin_elms += [Latex.arg(x) for x in args]\n\n env_elms: list[LatexLike] = [Latex.join(\"\", begin_elms)]\n if content is not None:\n env_elms.append(content)\n env_elms.append(rf\"\\end{{{name}}}\")\n\n return Latex.join(\" \", env_elms)\n\n @staticmethod\n def join(separator: LatexLike, elements: Iterable[LatexLike]) -> Latex:\n \"\"\"Joins given sequence.\n\n Args:\n separator: Expression of the separator between each element.\n elements: Iterable of expressions to be joined.\n\n Returns:\n A new Latex: \"{e[0]}{s}{e[1]}{s}...{s}{e[-1]}\"\n where s == separator, and e == elements.\n \"\"\"\n return Latex(str(separator).join(str(x) for x in elements))\n"
},
{
"path": "src/latexify/codegen/latex_test.py",
"content": "\"\"\"Tests for latexify.codegen.latex.\"\"\"\n\nfrom __future__ import annotations\n\n# Ignores [22-imports] for convenience.\nfrom latexify.codegen.latex import Latex\n\n\ndef test_eq() -> None:\n assert Latex(\"foo\") == Latex(\"foo\")\n assert Latex(\"foo\") != \"foo\"\n assert Latex(\"foo\") != Latex(\"bar\")\n\n\ndef test_str() -> None:\n assert str(Latex(\"foo\")) == \"foo\"\n\n\ndef test_add() -> None:\n assert Latex(\"foo\") + \"bar\" == Latex(\"foobar\")\n assert \"foo\" + Latex(\"bar\") == Latex(\"foobar\")\n assert Latex(\"foo\") + Latex(\"bar\") == Latex(\"foobar\")\n\n\ndef test_opt() -> None:\n assert Latex.opt(\"foo\") == Latex(\"[foo]\")\n assert Latex.opt(Latex(\"foo\")) == Latex(\"[foo]\")\n\n\ndef test_arg() -> None:\n assert Latex.arg(\"foo\") == Latex(\"{foo}\")\n assert Latex.arg(Latex(\"foo\")) == Latex(\"{foo}\")\n\n\ndef test_paren() -> None:\n assert Latex.paren(\"foo\") == Latex(r\"\\mathopen{}\\left( foo \\mathclose{}\\right)\")\n assert Latex.paren(Latex(\"foo\")) == Latex(\n r\"\\mathopen{}\\left( foo \\mathclose{}\\right)\"\n )\n\n\ndef test_curly() -> None:\n assert Latex.curly(\"foo\") == Latex(r\"\\mathopen{}\\left\\{ foo \\mathclose{}\\right\\}\")\n assert Latex.curly(Latex(\"foo\")) == Latex(\n r\"\\mathopen{}\\left\\{ foo \\mathclose{}\\right\\}\"\n )\n\n\ndef test_square() -> None:\n assert Latex.square(\"foo\") == Latex(r\"\\mathopen{}\\left[ foo \\mathclose{}\\right]\")\n assert Latex.square(Latex(\"foo\")) == Latex(\n r\"\\mathopen{}\\left[ foo \\mathclose{}\\right]\"\n )\n\n\ndef test_command() -> None:\n assert Latex.command(\"a\") == Latex(r\"\\a\")\n assert Latex.command(\"a\", options=[]) == Latex(r\"\\a\")\n assert Latex.command(\"a\", options=[\"b\"]) == Latex(r\"\\a[b]\")\n assert Latex.command(\"a\", options=[Latex(\"b\")]) == Latex(r\"\\a[b]\")\n assert Latex.command(\"a\", options=[\"b\", \"c\"]) == Latex(r\"\\a[b][c]\")\n assert Latex.command(\"a\", args=[]) == Latex(r\"\\a\")\n assert Latex.command(\"a\", args=[\"b\"]) == Latex(r\"\\a{b}\")\n assert Latex.command(\"a\", args=[Latex(\"b\")]) == Latex(r\"\\a{b}\")\n assert Latex.command(\"a\", args=[\"b\", \"c\"]) == Latex(r\"\\a{b}{c}\")\n assert Latex.command(\"a\", options=[\"b\"], args=[\"c\"]) == Latex(r\"\\a[b]{c}\")\n\n\ndef test_environment() -> None:\n assert Latex.environment(\"a\") == Latex(r\"\\begin{a} \\end{a}\")\n assert Latex.environment(\"a\", options=[]) == Latex(r\"\\begin{a} \\end{a}\")\n assert Latex.environment(\"a\", options=[\"b\"]) == Latex(r\"\\begin{a}[b] \\end{a}\")\n assert Latex.environment(\"a\", options=[Latex(\"b\")]) == Latex(\n r\"\\begin{a}[b] \\end{a}\"\n )\n assert Latex.environment(\"a\", options=[\"b\", \"c\"]) == Latex(\n r\"\\begin{a}[b][c] \\end{a}\"\n )\n assert Latex.environment(\"a\", args=[]) == Latex(r\"\\begin{a} \\end{a}\")\n assert Latex.environment(\"a\", args=[\"b\"]) == Latex(r\"\\begin{a}{b} \\end{a}\")\n assert Latex.environment(\"a\", args=[Latex(\"b\")]) == Latex(r\"\\begin{a}{b} \\end{a}\")\n assert Latex.environment(\"a\", args=[\"b\", \"c\"]) == Latex(r\"\\begin{a}{b}{c} \\end{a}\")\n assert Latex.environment(\"a\", content=\"b\") == Latex(r\"\\begin{a} b \\end{a}\")\n assert Latex.environment(\"a\", content=Latex(\"b\")) == Latex(r\"\\begin{a} b \\end{a}\")\n assert Latex.environment(\"a\", options=[\"b\"], args=[\"c\"]) == Latex(\n r\"\\begin{a}[b]{c} \\end{a}\"\n )\n assert Latex.environment(\"a\", options=[\"b\"], content=\"c\") == Latex(\n r\"\\begin{a}[b] c \\end{a}\"\n )\n assert Latex.environment(\"a\", args=[\"b\"], content=\"c\") == Latex(\n r\"\\begin{a}{b} c \\end{a}\"\n )\n assert Latex.environment(\"a\", options=[\"b\"], args=[\"c\"], content=\"d\") == Latex(\n r\"\\begin{a}[b]{c} d \\end{a}\"\n )\n\n\ndef test_join() -> None:\n assert Latex.join(\":\", []) == Latex(\"\")\n assert Latex.join(\":\", [\"foo\"]) == Latex(\"foo\")\n assert Latex.join(\":\", [Latex(\"foo\")]) == Latex(\"foo\")\n assert Latex.join(\":\", [Latex(\"foo\"), \"bar\"]) == Latex(\"foo:bar\")\n assert Latex.join(\":\", [\"foo\", Latex(\"bar\")]) == Latex(\"foo:bar\")\n assert Latex.join(\":\", [Latex(\"foo\"), Latex(\"bar\")]) == Latex(\"foo:bar\")\n assert Latex.join(\":\", ()) == Latex(\"\")\n assert Latex.join(\":\", (\"foo\",)) == Latex(\"foo\")\n assert Latex.join(\":\", (Latex(\"foo\"),)) == Latex(\"foo\")\n assert Latex.join(\":\", (Latex(\"foo\"), \"bar\")) == Latex(\"foo:bar\")\n assert Latex.join(\":\", (\"foo\", Latex(\"bar\"))) == Latex(\"foo:bar\")\n assert Latex.join(\":\", (Latex(\"foo\"), Latex(\"bar\"))) == Latex(\"foo:bar\")\n assert Latex.join(\":\", (str(x) for x in range(3))) == Latex(\"0:1:2\")\n assert Latex.join(\":\", (Latex(str(x)) for x in range(3))) == Latex(\"0:1:2\")\n"
},
{
"path": "src/latexify/config.py",
"content": "\"\"\"Definition of the Config class.\"\"\"\n\nfrom __future__ import annotations\n\nimport dataclasses\nfrom typing import Any\n\n\n@dataclasses.dataclass(frozen=True)\nclass Config:\n \"\"\"Configurations to control the behavior of latexify.\n\n Attributes:\n expand_functions: If set, the names of the functions to expand.\n identifiers: If set, the mapping to replace identifier names in the\n function. Keys are the original names of the identifiers,\n and corresponding values are the replacements.\n Both keys and values have to represent valid Python identifiers:\n ^[A-Za-z_][A-Za-z0-9_]*$\n prefixes: Prefixes of identifiers to trim. E.g., if \"foo.bar\" in prefixes, all\n identifiers with the form \"foo.bar.suffix\" will be replaced to \"suffix\"\n reduce_assignments: If True, assignment statements are used to synthesize\n the final expression.\n use_math_symbols: Whether to convert identifiers with a math symbol surface\n (e.g., \"alpha\") to the LaTeX symbol (e.g., \"\\\\alpha\").\n use_set_symbols: Whether to use set symbols or not.\n use_signature: Whether to add the function signature before the expression\n or not.\n \"\"\"\n\n expand_functions: set[str] | None\n identifiers: dict[str, str] | None\n prefixes: set[str] | None\n reduce_assignments: bool\n use_math_symbols: bool\n use_set_symbols: bool\n use_signature: bool\n escape_underscores: bool\n\n def merge(self, *, config: Config | None = None, **kwargs) -> Config:\n \"\"\"Merge configuration based on old configuration and field values.\n\n Args:\n config: If None, the merged one will merge defaults and field values,\n instead of merging old configuration and field values.\n **kwargs: Members to modify. This value precedes both self and config.\n\n Returns:\n A new Config object\n \"\"\"\n\n def merge_field(name: str) -> Any:\n # Precedence: kwargs -> config -> self\n arg = kwargs.get(name)\n if arg is None:\n if config is not None:\n arg = getattr(config, name)\n else:\n arg = getattr(self, name)\n return arg\n\n return Config(**{f.name: merge_field(f.name) for f in dataclasses.fields(self)})\n\n @staticmethod\n def defaults() -> Config:\n \"\"\"Generates a Config with default values.\n\n Returns:\n A new Config with default values\n \"\"\"\n return Config(\n expand_functions=None,\n identifiers=None,\n prefixes=None,\n reduce_assignments=False,\n use_math_symbols=False,\n use_set_symbols=False,\n use_signature=True,\n escape_underscores=True,\n )\n"
},
{
"path": "src/latexify/exceptions.py",
"content": "\"\"\"Exceptions used in Latexify.\"\"\"\n\n\nclass LatexifyError(Exception):\n \"\"\"Base class of all Latexify exceptions.\n\n Subclasses of this exception does not mean incorrect use of the library by the user\n at the interface level. These exceptions inform users that Latexify went into\n something wrong during processing the given functions.\n These exceptions are usually captured by the frontend functions (e.g., `with_latex`)\n to prevent destroying the entire program.\n Errors caused by wrong inputs should raise built-in exceptions.\n \"\"\"\n\n ...\n\n\nclass LatexifyNotSupportedError(LatexifyError):\n \"\"\"Some subtree in the AST is not supported by the current implementation.\n\n This error is raised when the library discovered incompatible syntaxes due to lack\n of the implementation. Possibly this error would be resolved in the future.\n \"\"\"\n\n ...\n\n\nclass LatexifySyntaxError(LatexifyError):\n \"\"\"Some subtree in the AST is not supported.\n\n This error is raised when the library discovered syntaxes that are not possible to\n be processed anymore. This error is essential, and wouldn't be resolved in the\n future.\n \"\"\"\n\n ...\n"
},
{
"path": "src/latexify/frontend.py",
"content": "\"\"\"Frontend interfaces of latexify.\"\"\"\n\nfrom __future__ import annotations\n\nfrom collections.abc import Callable\nfrom typing import Any, overload\n\nfrom latexify import ipython_wrappers\n\n\n@overload\ndef algorithmic(\n fn: Callable[..., Any], **kwargs: Any\n) -> ipython_wrappers.LatexifiedAlgorithm: ...\n\n\n@overload\ndef algorithmic(\n **kwargs: Any,\n) -> Callable[[Callable[..., Any]], ipython_wrappers.LatexifiedAlgorithm]: ...\n\n\ndef algorithmic(\n fn: Callable[..., Any] | None = None, **kwargs: Any\n) -> (\n ipython_wrappers.LatexifiedAlgorithm\n | Callable[[Callable[..., Any]], ipython_wrappers.LatexifiedAlgorithm]\n):\n \"\"\"Attach LaTeX pretty-printing to the given function.\n\n This function works with or without specifying the target function as the\n positional argument. The following two syntaxes works similarly.\n - latexify.algorithmic(alg, **kwargs)\n - latexify.algorithmic(**kwargs)(alg)\n\n Args:\n fn: Callable to be wrapped.\n **kwargs: Arguments to control behavior. See also get_latex().\n\n Returns:\n - If `fn` is passed, returns the wrapped function.\n - Otherwise, returns the wrapper function with given settings.\n \"\"\"\n if fn is not None:\n return ipython_wrappers.LatexifiedAlgorithm(fn, **kwargs)\n\n def wrapper(f):\n return ipython_wrappers.LatexifiedAlgorithm(f, **kwargs)\n\n return wrapper\n\n\n@overload\ndef function(\n fn: Callable[..., Any], **kwargs: Any\n) -> ipython_wrappers.LatexifiedFunction: ...\n\n\n@overload\ndef function(\n **kwargs: Any,\n) -> Callable[[Callable[..., Any]], ipython_wrappers.LatexifiedFunction]: ...\n\n\ndef function(\n fn: Callable[..., Any] | None = None, **kwargs: Any\n) -> (\n ipython_wrappers.LatexifiedFunction\n | Callable[[Callable[..., Any]], ipython_wrappers.LatexifiedFunction]\n):\n \"\"\"Attach LaTeX pretty-printing to the given function.\n\n This function works with or without specifying the target function as the positional\n argument. The following two syntaxes works similarly.\n - latexify.function(fn, **kwargs)\n - latexify.function(**kwargs)(fn)\n\n Args:\n fn: Callable to be wrapped.\n **kwargs: Arguments to control behavior. See also get_latex().\n\n Returns:\n - If `fn` is passed, returns the wrapped function.\n - Otherwise, returns the wrapper function with given settings.\n \"\"\"\n if fn is not None:\n return ipython_wrappers.LatexifiedFunction(fn, **kwargs)\n\n def wrapper(f):\n return ipython_wrappers.LatexifiedFunction(f, **kwargs)\n\n return wrapper\n\n\n@overload\ndef expression(\n fn: Callable[..., Any], **kwargs: Any\n) -> ipython_wrappers.LatexifiedFunction: ...\n\n\n@overload\ndef expression(\n **kwargs: Any,\n) -> Callable[[Callable[..., Any]], ipython_wrappers.LatexifiedFunction]: ...\n\n\ndef expression(\n fn: Callable[..., Any] | None = None, **kwargs: Any\n) -> (\n ipython_wrappers.LatexifiedFunction\n | Callable[[Callable[..., Any]], ipython_wrappers.LatexifiedFunction]\n):\n \"\"\"Attach LaTeX pretty-printing to the given function.\n\n This function is a shortcut for `latexify.function` with the default parameter\n `use_signature=False`.\n \"\"\"\n kwargs[\"use_signature\"] = kwargs.get(\"use_signature\", False)\n\n if fn is not None:\n return ipython_wrappers.LatexifiedFunction(fn, **kwargs)\n\n def wrapper(f):\n return ipython_wrappers.LatexifiedFunction(f, **kwargs)\n\n return wrapper\n"
},
{
"path": "src/latexify/frontend_test.py",
"content": "\"\"\"Tests for latexify.frontend.\"\"\"\n\nfrom __future__ import annotations\n\nfrom latexify import frontend\n\n\ndef test_function() -> None:\n def f(x):\n return x\n\n latex_without_flag = \"x\"\n latex_with_flag = r\"f(x) = x\"\n\n # Checks the syntax:\n # @function\n # def fn(...):\n # ...\n latexified = frontend.function(f)\n assert str(latexified) == latex_with_flag\n assert latexified._repr_latex_() == rf\"$$ \\displaystyle {latex_with_flag} $$\"\n\n # Checks the syntax:\n # @function(**kwargs)\n # def fn(...):\n # ...\n latexified = frontend.function(use_signature=False)(f)\n assert str(latexified) == latex_without_flag\n assert latexified._repr_latex_() == rf\"$$ \\displaystyle {latex_without_flag} $$\"\n\n # Checks the syntax:\n # def fn(...):\n # ...\n # latexified = function(fn, **kwargs)\n latexified = frontend.function(f, use_signature=False)\n assert str(latexified) == latex_without_flag\n assert latexified._repr_latex_() == rf\"$$ \\displaystyle {latex_without_flag} $$\"\n\n\ndef test_expression() -> None:\n def f(x):\n return x\n\n latex_without_flag = \"x\"\n latex_with_flag = r\"f(x) = x\"\n\n # Checks the syntax:\n # @expression\n # def fn(...):\n # ...\n latexified = frontend.expression(f)\n assert str(latexified) == latex_without_flag\n assert latexified._repr_latex_() == rf\"$$ \\displaystyle {latex_without_flag} $$\"\n\n # Checks the syntax:\n # @expression(**kwargs)\n # def fn(...):\n # ...\n latexified = frontend.expression(use_signature=True)(f)\n assert str(latexified) == latex_with_flag\n assert latexified._repr_latex_() == rf\"$$ \\displaystyle {latex_with_flag} $$\"\n\n # Checks the syntax:\n # def fn(...):\n # ...\n # latexified = expression(fn, **kwargs)\n latexified = frontend.expression(f, use_signature=True)\n assert str(latexified) == latex_with_flag\n assert latexified._repr_latex_() == rf\"$$ \\displaystyle {latex_with_flag} $$\"\n"
},
{
"path": "src/latexify/generate_latex.py",
"content": "\"\"\"Generate LaTeX code.\"\"\"\n\nfrom __future__ import annotations\n\nimport enum\nfrom collections.abc import Callable\nfrom typing import Any\n\nfrom latexify import codegen\nfrom latexify import config as cfg\nfrom latexify import parser, transformers\n\n\nclass Style(enum.Enum):\n \"\"\"The style of the generated LaTeX.\"\"\"\n\n ALGORITHMIC = \"algorithmic\"\n FUNCTION = \"function\"\n IPYTHON_ALGORITHMIC = \"ipython-algorithmic\"\n\n\ndef get_latex(\n fn: Callable[..., Any],\n *,\n style: Style = Style.FUNCTION,\n config: cfg.Config | None = None,\n **kwargs,\n) -> str:\n \"\"\"Obtains LaTeX description from the function's source.\n\n Args:\n fn: Reference to a function to analyze.\n style: Style of the LaTeX description, the default is FUNCTION.\n config: Use defined Config object, if it is None, it will be automatic assigned\n with default value.\n **kwargs: Dict of Config field values that could be defined individually\n by users.\n\n Returns:\n Generated LaTeX description.\n\n Raises:\n latexify.exceptions.LatexifyError: Something went wrong during conversion.\n \"\"\"\n merged_config = cfg.Config.defaults().merge(config=config, **kwargs)\n\n # Obtains the source AST.\n tree = parser.parse_function(fn)\n\n # Mandatory AST Transformation.\n tree = transformers.AugAssignReplacer().visit(tree)\n\n # Conditional AST transformation.\n if merged_config.prefixes is not None:\n tree = transformers.PrefixTrimmer(merged_config.prefixes).visit(tree)\n if merged_config.identifiers is not None:\n tree = transformers.IdentifierReplacer(merged_config.identifiers).visit(tree)\n if merged_config.reduce_assignments:\n tree = transformers.DocstringRemover().visit(tree)\n tree = transformers.AssignmentReducer().visit(tree)\n if merged_config.expand_functions is not None:\n tree = transformers.FunctionExpander(merged_config.expand_functions).visit(tree)\n\n # Generates LaTeX.\n if style == Style.ALGORITHMIC:\n return codegen.AlgorithmicCodegen(\n use_math_symbols=merged_config.use_math_symbols,\n use_set_symbols=merged_config.use_set_symbols,\n escape_underscores=merged_config.escape_underscores,\n ).visit(tree)\n elif style == Style.FUNCTION:\n return codegen.FunctionCodegen(\n use_math_symbols=merged_config.use_math_symbols,\n use_signature=merged_config.use_signature,\n use_set_symbols=merged_config.use_set_symbols,\n escape_underscores=merged_config.escape_underscores,\n ).visit(tree)\n elif style == Style.IPYTHON_ALGORITHMIC:\n return codegen.IPythonAlgorithmicCodegen(\n use_math_symbols=merged_config.use_math_symbols,\n use_set_symbols=merged_config.use_set_symbols,\n escape_underscores=merged_config.escape_underscores,\n ).visit(tree)\n\n raise ValueError(f\"Unrecognized style: {style}\")\n"
},
{
"path": "src/latexify/generate_latex_test.py",
"content": "\"\"\"Tests for latexify.generate_latex.\"\"\"\n\nfrom __future__ import annotations\n\nfrom latexify import generate_latex\n\n\ndef test_get_latex_identifiers() -> None:\n def myfn(myvar):\n return 3 * myvar\n\n identifiers = {\"myfn\": \"f\", \"myvar\": \"x\"}\n\n latex_without_flag = r\"\\mathrm{myfn}(\\mathrm{myvar}) = 3 \\mathrm{myvar}\"\n latex_with_flag = r\"f(x) = 3 x\"\n\n assert generate_latex.get_latex(myfn) == latex_without_flag\n assert generate_latex.get_latex(myfn, identifiers=identifiers) == latex_with_flag\n\n\ndef test_get_latex_prefixes() -> None:\n abc = object()\n\n def f(x):\n return abc.d + x.y.z.e\n\n latex_without_flag = r\"f(x) = \\mathrm{abc}.d + x.y.z.e\"\n latex_with_flag1 = r\"f(x) = d + x.y.z.e\"\n latex_with_flag2 = r\"f(x) = \\mathrm{abc}.d + y.z.e\"\n latex_with_flag3 = r\"f(x) = \\mathrm{abc}.d + z.e\"\n latex_with_flag4 = r\"f(x) = d + e\"\n\n assert generate_latex.get_latex(f) == latex_without_flag\n assert generate_latex.get_latex(f, prefixes=set()) == latex_without_flag\n assert generate_latex.get_latex(f, prefixes={\"abc\"}) == latex_with_flag1\n assert generate_latex.get_latex(f, prefixes={\"x\"}) == latex_with_flag2\n assert generate_latex.get_latex(f, prefixes={\"x.y\"}) == latex_with_flag3\n assert generate_latex.get_latex(f, prefixes={\"abc\", \"x.y.z\"}) == latex_with_flag4\n assert (\n generate_latex.get_latex(f, prefixes={\"abc\", \"x\", \"x.y.z\"}) == latex_with_flag4\n )\n\n\ndef test_get_latex_reduce_assignments() -> None:\n def f(x):\n y = 3 * x\n return y\n\n latex_without_flag = r\"\\begin{array}{l} y = 3 x \\\\ f(x) = y \\end{array}\"\n latex_with_flag = r\"f(x) = 3 x\"\n\n assert generate_latex.get_latex(f) == latex_without_flag\n assert generate_latex.get_latex(f, reduce_assignments=False) == latex_without_flag\n assert generate_latex.get_latex(f, reduce_assignments=True) == latex_with_flag\n\n\ndef test_get_latex_reduce_assignments_with_docstring() -> None:\n def f(x):\n \"\"\"DocstringRemover is required.\"\"\"\n y = 3 * x\n return y\n\n latex_without_flag = r\"\\begin{array}{l} y = 3 x \\\\ f(x) = y \\end{array}\"\n latex_with_flag = r\"f(x) = 3 x\"\n\n assert generate_latex.get_latex(f) == latex_without_flag\n assert generate_latex.get_latex(f, reduce_assignments=False) == latex_without_flag\n assert generate_latex.get_latex(f, reduce_assignments=True) == latex_with_flag\n\n\ndef test_get_latex_reduce_assignments_with_aug_assign() -> None:\n def f(x):\n y = 3\n y *= x\n return y\n\n latex_without_flag = r\"\\begin{array}{l} y = 3 \\\\ y = y x \\\\ f(x) = y \\end{array}\"\n latex_with_flag = r\"f(x) = 3 x\"\n\n assert generate_latex.get_latex(f) == latex_without_flag\n assert generate_latex.get_latex(f, reduce_assignments=False) == latex_without_flag\n assert generate_latex.get_latex(f, reduce_assignments=True) == latex_with_flag\n\n\ndef test_get_latex_use_math_symbols() -> None:\n def f(alpha):\n return alpha\n\n latex_without_flag = r\"f(\\mathrm{alpha}) = \\mathrm{alpha}\"\n latex_with_flag = r\"f(\\alpha) = \\alpha\"\n\n assert generate_latex.get_latex(f) == latex_without_flag\n assert generate_latex.get_latex(f, use_math_symbols=False) == latex_without_flag\n assert generate_latex.get_latex(f, use_math_symbols=True) == latex_with_flag\n\n\ndef test_get_latex_use_signature() -> None:\n def f(x):\n return x\n\n latex_without_flag = \"x\"\n latex_with_flag = r\"f(x) = x\"\n\n assert generate_latex.get_latex(f) == latex_with_flag\n assert generate_latex.get_latex(f, use_signature=False) == latex_without_flag\n assert generate_latex.get_latex(f, use_signature=True) == latex_with_flag\n\n\ndef test_get_latex_use_set_symbols() -> None:\n def f(x, y):\n return x & y\n\n latex_without_flag = r\"f(x, y) = x \\mathbin{\\&} y\"\n latex_with_flag = r\"f(x, y) = x \\cap y\"\n\n assert generate_latex.get_latex(f) == latex_without_flag\n assert generate_latex.get_latex(f, use_set_symbols=False) == latex_without_flag\n assert generate_latex.get_latex(f, use_set_symbols=True) == latex_with_flag\n"
},
{
"path": "src/latexify/ipython_wrappers.py",
"content": "\"\"\"Wrapper objects for IPython to display output.\"\"\"\n\nfrom __future__ import annotations\n\nimport abc\nfrom typing import Any, Callable, cast\n\nfrom latexify import exceptions, generate_latex\n\n\nclass LatexifiedRepr(metaclass=abc.ABCMeta):\n \"\"\"Object with LaTeX representation.\"\"\"\n\n _fn: Callable[..., Any]\n\n def __init__(self, fn: Callable[..., Any], **kwargs) -> None:\n self._fn = fn\n\n @property\n def __doc__(self) -> str | None:\n return self._fn.__doc__\n\n @__doc__.setter\n def __doc__(self, val: str | None) -> None:\n self._fn.__doc__ = val\n\n @property\n def __name__(self) -> str:\n return self._fn.__name__\n\n @__name__.setter\n def __name__(self, val: str) -> None:\n self._fn.__name__ = val\n\n # After Python 3.7\n # @final\n def __call__(self, *args) -> Any:\n return self._fn(*args)\n\n @abc.abstractmethod\n def __str__(self) -> str: ...\n\n @abc.abstractmethod\n def _repr_html_(self) -> str | tuple[str, dict[str, Any]] | None:\n \"\"\"IPython hook to display HTML visualization.\"\"\"\n ...\n\n @abc.abstractmethod\n def _repr_latex_(self) -> str | tuple[str, dict[str, Any]] | None:\n \"\"\"IPython hook to display LaTeX visualization.\"\"\"\n ...\n\n\nclass LatexifiedAlgorithm(LatexifiedRepr):\n \"\"\"Algorithm with latex representation.\"\"\"\n\n _latex: str | None\n _error: str | None\n _ipython_latex: str | None\n _ipython_error: str | None\n\n def __init__(self, fn: Callable[..., Any], **kwargs) -> None:\n super().__init__(fn)\n\n try:\n self._latex = generate_latex.get_latex(\n fn, style=generate_latex.Style.ALGORITHMIC, **kwargs\n )\n self._error = None\n except exceptions.LatexifyError as e:\n self._latex = None\n self._error = f\"{type(e).__name__}: {str(e)}\"\n\n try:\n self._ipython_latex = generate_latex.get_latex(\n fn, style=generate_latex.Style.IPYTHON_ALGORITHMIC, **kwargs\n )\n self._ipython_error = None\n except exceptions.LatexifyError as e:\n self._ipython_latex = None\n self._ipython_error = f\"{type(e).__name__}: {str(e)}\"\n\n def __str__(self) -> str:\n return self._latex if self._latex is not None else cast(str, self._error)\n\n def _repr_html_(self) -> str | tuple[str, dict[str, Any]] | None:\n \"\"\"IPython hook to display HTML visualization.\"\"\"\n return (\n '' + self._ipython_error + \"\"\n if self._ipython_error is not None\n else None\n )\n\n def _repr_latex_(self) -> str | tuple[str, dict[str, Any]] | None:\n \"\"\"IPython hook to display LaTeX visualization.\"\"\"\n return (\n f\"$ {self._ipython_latex} $\"\n if self._ipython_latex is not None\n else self._ipython_error\n )\n\n\nclass LatexifiedFunction(LatexifiedRepr):\n \"\"\"Function with latex representation.\"\"\"\n\n _latex: str | None\n _error: str | None\n\n def __init__(self, fn: Callable[..., Any], **kwargs) -> None:\n super().__init__(fn, **kwargs)\n\n try:\n self._latex = self._latex = generate_latex.get_latex(\n fn, style=generate_latex.Style.FUNCTION, **kwargs\n )\n self._error = None\n except exceptions.LatexifyError as e:\n self._latex = None\n self._error = f\"{type(e).__name__}: {str(e)}\"\n\n def __str__(self) -> str:\n return self._latex if self._latex is not None else cast(str, self._error)\n\n def _repr_html_(self) -> str | tuple[str, dict[str, Any]] | None:\n \"\"\"IPython hook to display HTML visualization.\"\"\"\n return (\n '' + self._error + \"\"\n if self._error is not None\n else None\n )\n\n def _repr_latex_(self) -> str | tuple[str, dict[str, Any]] | None:\n \"\"\"IPython hook to display LaTeX visualization.\"\"\"\n return (\n rf\"$$ \\displaystyle {self._latex} $$\"\n if self._latex is not None\n else self._error\n )\n"
},
{
"path": "src/latexify/parser.py",
"content": "\"\"\"Parsing utilities.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\nimport inspect\nimport textwrap\nfrom collections.abc import Callable\nfrom typing import Any\n\nimport dill # type: ignore[import]\n\nfrom latexify import exceptions\n\n\ndef parse_function(fn: Callable[..., Any]) -> ast.Module:\n \"\"\"Parses given function.\n\n Args:\n fn: Target function.\n\n Returns:\n AST tree representing `fn`.\n \"\"\"\n try:\n source = inspect.getsource(fn)\n except Exception:\n # Maybe running on console.\n source = dill.source.getsource(fn)\n\n # Remove extra indentation so that ast.parse runs correctly.\n source = textwrap.dedent(source)\n\n tree = ast.parse(source)\n if not tree.body or not isinstance(tree.body[0], ast.FunctionDef):\n raise exceptions.LatexifySyntaxError(\"Not a function.\")\n\n return tree\n"
},
{
"path": "src/latexify/parser_test.py",
"content": "\"\"\"Tests for latexify.parser.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\n\nimport pytest\n\nfrom latexify import ast_utils, exceptions, parser, test_utils\n\n\ndef test_parse_function_with_posonlyargs() -> None:\n def f(x):\n return x\n\n expected = ast.Module(\n body=[\n ast_utils.create_function_def(\n name=\"f\",\n args=ast.arguments(\n posonlyargs=[],\n args=[ast.arg(arg=\"x\")],\n vararg=None,\n kwonlyargs=[],\n kw_defaults=[],\n kwarg=None,\n defaults=[],\n ),\n body=[ast.Return(value=ast.Name(id=\"x\", ctx=ast.Load()))],\n decorator_list=[],\n returns=None,\n type_comment=None,\n type_params=[],\n lineno=1,\n col_offset=0,\n end_lineno=2,\n end_col_offset=0,\n )\n ],\n type_ignores=[],\n )\n\n obtained = parser.parse_function(f)\n test_utils.assert_ast_equal(obtained, expected)\n\n\ndef test_parse_function_with_lambda() -> None:\n with pytest.raises(exceptions.LatexifySyntaxError, match=r\"^Not a function\\.$\"):\n parser.parse_function(lambda: ())\n with pytest.raises(exceptions.LatexifySyntaxError, match=r\"^Not a function\\.$\"):\n x = lambda: () # noqa: E731\n parser.parse_function(x)\n"
},
{
"path": "src/latexify/test_utils.py",
"content": "\"\"\"Test utilities.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\nimport functools\nimport sys\nfrom collections.abc import Callable\nfrom typing import cast\n\n\ndef require_at_least(\n minor: int,\n) -> Callable[[Callable[..., None]], Callable[..., None]]:\n \"\"\"Require the minimum minor version of Python 3 to run the test.\n\n Args:\n minor: Minimum minor version (inclusive) that the test case supports.\n\n Returns:\n A decorator function to wrap the test case function.\n \"\"\"\n\n def decorator(fn: Callable[..., None]) -> Callable[..., None]:\n @functools.wraps(fn)\n def wrapper(*args, **kwargs):\n if sys.version_info.minor < minor:\n return\n fn(*args, **kwargs)\n\n return wrapper\n\n return decorator\n\n\ndef require_at_most(\n minor: int,\n) -> Callable[[Callable[..., None]], Callable[..., None]]:\n \"\"\"Require the maximum minor version of Python 3 to run the test.\n\n Args:\n minor: Maximum minor version (inclusive) that the test case supports.\n\n Returns:\n A decorator function to wrap the test case function.\n \"\"\"\n\n def decorator(fn: Callable[..., None]) -> Callable[..., None]:\n @functools.wraps(fn)\n def wrapper(*args, **kwargs):\n if sys.version_info.minor > minor:\n return\n fn(*args, **kwargs)\n\n return wrapper\n\n return decorator\n\n\ndef ast_equal(observed: ast.AST, expected: ast.AST) -> bool:\n \"\"\"Checks the equality between two ASTs.\n\n This function checks if `observed` contains at least the same subtree with\n `expected`. If `observed` has some extra branches that `expected` does not cover,\n it is ignored.\n\n Args:\n observed: An AST to check.\n expected: The expected AST.\n\n Returns:\n True if observed and expected represent the same AST, False otherwise.\n \"\"\"\n ignore_keys = {\"lineno\", \"col_offset\", \"end_lineno\", \"end_col_offset\", \"kind\"}\n if sys.version_info.minor <= 12:\n ignore_keys.add(\"type_params\")\n\n try:\n assert type(observed) is type(expected)\n\n for k, ve in vars(expected).items():\n if k in ignore_keys:\n continue\n\n vo = getattr(observed, k) # May cause AttributeError.\n\n if isinstance(ve, ast.AST):\n assert ast_equal(cast(ast.AST, vo), ve)\n elif isinstance(ve, list):\n vo = cast(list, vo)\n assert len(vo) == len(ve)\n assert all(\n ast_equal(cast(ast.AST, co), cast(ast.AST, ce))\n for co, ce in zip(vo, ve)\n )\n else:\n assert type(vo) is type(ve)\n assert vo == ve\n\n except (AssertionError, AttributeError):\n raise # raise to debug easier.\n\n return True\n\n\ndef assert_ast_equal(observed: ast.AST, expected: ast.AST) -> None:\n \"\"\"Asserts the equality between two ASTs.\n\n Args:\n observed: An AST to compare.\n expected: Another AST.\n\n Raises:\n AssertionError: observed and expected represent different ASTs.\n \"\"\"\n assert ast_equal(\n observed, expected\n ), f\"\"\"\\\n AST does not match.\n observed={ast.dump(observed, indent=4)}\n expected={ast.dump(expected, indent=4)}\n \"\"\"\n"
},
{
"path": "src/latexify/transformers/__init__.py",
"content": "\"\"\"Package latexify.transformers.\"\"\"\n\nfrom latexify.transformers.assignment_reducer import AssignmentReducer\nfrom latexify.transformers.aug_assign_replacer import AugAssignReplacer\nfrom latexify.transformers.docstring_remover import DocstringRemover\nfrom latexify.transformers.function_expander import FunctionExpander\nfrom latexify.transformers.identifier_replacer import IdentifierReplacer\nfrom latexify.transformers.prefix_trimmer import PrefixTrimmer\n\n__all__ = [\n \"AssignmentReducer\",\n \"AugAssignReplacer\",\n \"DocstringRemover\",\n \"FunctionExpander\",\n \"IdentifierReplacer\",\n \"PrefixTrimmer\",\n]\n"
},
{
"path": "src/latexify/transformers/assignment_reducer.py",
"content": "\"\"\"NodeTransformer to reduce assigned expressions.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\nfrom typing import Any\n\nfrom latexify import ast_utils, exceptions\n\n\nclass AssignmentReducer(ast.NodeTransformer):\n \"\"\"NodeTransformer to reduce assigned expressions.\n\n This class replaces a functions with multiple assignments to a function with only\n single return.\n\n Example:\n def f(x):\n y = 2 + x\n z = 3 * y\n return 4 + z\n\n AssignmentReducer modifies the function above to below:\n\n def f(x):\n return 4 + 3 * (2 + x)\n \"\"\"\n\n _assignments: dict[str, ast.expr] | None = None\n\n # TODO(odashi):\n # Currently, this function does not care much about some expressions, e.g.,\n # comprehensions or lambdas, which introduces inner scopes.\n # It may cause some mistakes in the resulting AST.\n def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:\n \"\"\"Visit a FunctionDef node.\"\"\"\n # Push stack\n parent_assignments = self._assignments\n self._assignments = {}\n\n for child in node.body[:-1]:\n if not isinstance(child, ast.Assign):\n raise exceptions.LatexifyNotSupportedError(\n \"AssignmentReducer supports only Assign nodes, \"\n f\"but got: {type(child).__name__}\"\n )\n\n value = self.visit(child.value)\n\n for target in child.targets:\n if not isinstance(target, ast.Name):\n raise exceptions.LatexifyNotSupportedError(\n \"AssignmentReducer does not recognize list/tuple \"\n \"decomposition.\"\n )\n self._assignments[target.id] = value\n\n return_original = node.body[-1]\n\n if not isinstance(return_original, (ast.Return, ast.If)):\n raise exceptions.LatexifySyntaxError(\n f\"Unsupported last statement: {type(return_original).__name__}\"\n )\n\n return_transformed = self.visit(return_original)\n\n # Pop stack\n self._assignments = parent_assignments\n type_params = getattr(node, \"type_params\", [])\n return ast_utils.create_function_def(\n name=node.name,\n args=node.args,\n body=[return_transformed],\n decorator_list=node.decorator_list,\n returns=node.returns,\n type_params=type_params,\n )\n\n def visit_Name(self, node: ast.Name) -> Any:\n \"\"\"Visit a Name node.\"\"\"\n if self._assignments is not None:\n return self._assignments.get(node.id, node)\n\n return node\n"
},
{
"path": "src/latexify/transformers/assignment_reducer_test.py",
"content": "\"\"\"Tests for latexify.transformers.assignment_reducer.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\n\nfrom latexify import ast_utils, parser, test_utils\nfrom latexify.transformers.assignment_reducer import AssignmentReducer\n\n\ndef _make_ast(body: list[ast.stmt]) -> ast.Module:\n \"\"\"Helper function to generate an AST for f(x).\n\n Args:\n body: The function body.\n\n Returns:\n Generated AST.\n \"\"\"\n return ast.Module(\n body=[\n ast_utils.create_function_def(\n name=\"f\",\n args=ast.arguments(\n args=[ast.arg(arg=\"x\")],\n kwonlyargs=[],\n kw_defaults=[],\n defaults=[],\n posonlyargs=[],\n ),\n body=body,\n decorator_list=[],\n type_params=[],\n )\n ],\n type_ignores=[],\n )\n\n\ndef test_unchanged() -> None:\n def f(x):\n return x\n\n expected = _make_ast(\n [\n ast.Return(value=ast.Name(id=\"x\", ctx=ast.Load())),\n ]\n )\n transformed = AssignmentReducer().visit(parser.parse_function(f))\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_constant() -> None:\n def f(x):\n y = 0\n return y\n\n expected = _make_ast(\n [\n ast.Return(value=ast_utils.make_constant(0)),\n ]\n )\n transformed = AssignmentReducer().visit(parser.parse_function(f))\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_nested() -> None:\n def f(x):\n y = 2 * x\n return y\n\n expected = _make_ast(\n [\n ast.Return(\n value=ast.BinOp(\n left=ast_utils.make_constant(2),\n op=ast.Mult(),\n right=ast.Name(id=\"x\", ctx=ast.Load()),\n )\n )\n ]\n )\n transformed = AssignmentReducer().visit(parser.parse_function(f))\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_nested2() -> None:\n def f(x):\n y = 2 * x\n z = 3 + y\n return z\n\n expected = _make_ast(\n [\n ast.Return(\n value=ast.BinOp(\n left=ast_utils.make_constant(3),\n op=ast.Add(),\n right=ast.BinOp(\n left=ast_utils.make_constant(2),\n op=ast.Mult(),\n right=ast.Name(id=\"x\", ctx=ast.Load()),\n ),\n )\n )\n ]\n )\n transformed = AssignmentReducer().visit(parser.parse_function(f))\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_overwrite() -> None:\n def f(x):\n y = 2 * x\n y = 3 + x\n return y\n\n expected = _make_ast(\n [\n ast.Return(\n value=ast.BinOp(\n left=ast_utils.make_constant(3),\n op=ast.Add(),\n right=ast.Name(id=\"x\", ctx=ast.Load()),\n )\n )\n ]\n )\n transformed = AssignmentReducer().visit(parser.parse_function(f))\n test_utils.assert_ast_equal(transformed, expected)\n"
},
{
"path": "src/latexify/transformers/aug_assign_replacer.py",
"content": "\"\"\"Transformer to replace AugAssign to Assign.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\n\n\nclass AugAssignReplacer(ast.NodeTransformer):\n \"\"\"NodeTransformer to replace AugAssign to corresponding Assign.\n\n AugAssign(target, op, value) => Assign([target], BinOp(target, op, value))\n\n \"\"\"\n\n def visit_AugAssign(self, node: ast.AugAssign) -> ast.Assign:\n left_args = {**vars(node.target), \"ctx\": ast.Load()}\n left = type(node.target)(**left_args)\n return ast.Assign(\n targets=[node.target], value=ast.BinOp(left, node.op, node.value)\n )\n"
},
{
"path": "src/latexify/transformers/aug_assign_replacer_test.py",
"content": "\"\"\"Tests for latexify.transformers.aug_assign_replacer.\"\"\"\n\nimport ast\n\nfrom latexify import test_utils\nfrom latexify.transformers.aug_assign_replacer import AugAssignReplacer\n\n\ndef test_replace() -> None:\n tree = ast.AugAssign(\n target=ast.Name(id=\"x\", ctx=ast.Store()),\n op=ast.Add(),\n value=ast.Name(id=\"y\", ctx=ast.Load()),\n )\n expected = ast.Assign(\n targets=[ast.Name(id=\"x\", ctx=ast.Store())],\n value=ast.BinOp(\n left=ast.Name(id=\"x\", ctx=ast.Load()),\n op=ast.Add(),\n right=ast.Name(id=\"y\", ctx=ast.Load()),\n ),\n )\n transformed = AugAssignReplacer().visit(tree)\n test_utils.assert_ast_equal(transformed, expected)\n"
},
{
"path": "src/latexify/transformers/docstring_remover.py",
"content": "\"\"\"Transformer to remove all docstrings.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\nfrom typing import Union\n\nfrom latexify import ast_utils\n\n\nclass DocstringRemover(ast.NodeTransformer):\n \"\"\"NodeTransformer to remove all docstrings.\n\n Docstrings here are detected as Expr nodes with a single string constant.\n \"\"\"\n\n def visit_Expr(self, node: ast.Expr) -> Union[ast.Expr, None]:\n if ast_utils.is_str(node.value):\n return None\n return node\n"
},
{
"path": "src/latexify/transformers/docstring_remover_test.py",
"content": "\"\"\"Tests for latexify.transformers.docstring_remover.\"\"\"\n\nimport ast\n\nfrom latexify import ast_utils, parser, test_utils\nfrom latexify.transformers.docstring_remover import DocstringRemover\n\n\ndef test_remove_docstrings() -> None:\n def f():\n \"\"\"Test docstring.\"\"\"\n x = 42\n f() # This Expr should not be removed.\n \"\"\"This string constant should also be removed.\"\"\"\n return x\n\n tree = parser.parse_function(f).body[0]\n assert isinstance(tree, ast.FunctionDef)\n\n expected = ast_utils.create_function_def(\n name=\"f\",\n body=[\n ast.Assign(\n targets=[ast.Name(id=\"x\", ctx=ast.Store())],\n value=ast_utils.make_constant(42),\n ),\n ast.Expr(\n value=ast.Call(\n func=ast.Name(id=\"f\", ctx=ast.Load()), args=[], keywords=[]\n )\n ),\n ast.Return(value=ast.Name(id=\"x\", ctx=ast.Load())),\n ],\n args=ast.arguments(\n posonlyargs=[],\n args=[],\n vararg=None,\n kwonlyargs=[],\n kw_defaults=[],\n kwarg=None,\n defaults=[],\n ),\n decorator_list=[],\n type_params=[],\n )\n transformed = DocstringRemover().visit(tree)\n test_utils.assert_ast_equal(transformed, expected)\n"
},
{
"path": "src/latexify/transformers/function_expander.py",
"content": "from __future__ import annotations\n\nimport ast\nimport functools\nfrom collections.abc import Callable\n\nfrom latexify import ast_utils, exceptions\n\n\n# TODO(ZibingZhang): handle mutually recursive function expansions\nclass FunctionExpander(ast.NodeTransformer):\n \"\"\"NodeTransformer to expand functions.\n\n This class replaces function calls with an expanded form.\n\n Example:\n def f(x, y):\n return hypot(x, y)\n\n FunctionExpander({\"hypot\"}) will modify the AST of the function above to below:\n\n def f(x, y):\n return sqrt(x**2, y**2)\n \"\"\"\n\n def __init__(self, functions: set[str]) -> None:\n self._functions = functions\n\n def visit_Call(self, node: ast.Call) -> ast.AST:\n \"\"\"Visit a Call node.\"\"\"\n func_name = ast_utils.extract_function_name_or_none(node)\n if (\n func_name is not None\n and func_name in self._functions\n and func_name in _FUNCTION_EXPANDERS\n ):\n return _FUNCTION_EXPANDERS[func_name](self, node)\n\n kwargs = {\n \"func\": self.visit(node.func),\n \"args\": [self.visit(x) for x in node.args],\n }\n\n if hasattr(node, \"keywords\"):\n kwargs[\"keywords\"] = [\n ast.keyword(arg=x.arg, value=self.visit(x.value)) for x in node.keywords\n ]\n\n return ast.Call(**kwargs)\n\n\ndef _atan2_expander(function_expander: FunctionExpander, node: ast.Call) -> ast.AST:\n _check_num_args(node, 2)\n return ast.Call(\n func=ast.Name(id=\"atan\", ctx=ast.Load()),\n args=[\n ast.BinOp(\n left=function_expander.visit(node.args[0]),\n op=ast.Div(),\n right=function_expander.visit(node.args[1]),\n )\n ],\n keywords=[],\n )\n\n\ndef _exp_expander(function_expander: FunctionExpander, node: ast.Call) -> ast.AST:\n _check_num_args(node, 1)\n return ast.BinOp(\n left=ast.Name(id=\"e\", ctx=ast.Load()),\n op=ast.Pow(),\n right=function_expander.visit(node.args[0]),\n )\n\n\ndef _exp2_expander(function_expander: FunctionExpander, node: ast.Call) -> ast.AST:\n _check_num_args(node, 1)\n return ast.BinOp(\n left=ast_utils.make_constant(2),\n op=ast.Pow(),\n right=function_expander.visit(node.args[0]),\n )\n\n\ndef _expm1_expander(function_expander: FunctionExpander, node: ast.Call) -> ast.AST:\n _check_num_args(node, 1)\n return ast.BinOp(\n left=function_expander.visit(\n ast.Call(\n func=ast.Name(id=\"exp\", ctx=ast.Load()),\n args=[node.args[0]],\n keywords=[],\n )\n ),\n op=ast.Sub(),\n right=ast_utils.make_constant(1),\n )\n\n\ndef _hypot_expander(function_expander: FunctionExpander, node: ast.Call) -> ast.AST:\n if not node.args:\n return ast_utils.make_constant(0)\n\n args = [\n ast.BinOp(\n left=function_expander.visit(arg),\n op=ast.Pow(),\n right=ast_utils.make_constant(2),\n )\n for arg in node.args\n ]\n\n args_reduced = functools.reduce(\n lambda a, b: ast.BinOp(left=a, op=ast.Add(), right=b), args\n )\n return ast.Call(\n func=ast.Name(id=\"sqrt\", ctx=ast.Load()),\n args=[args_reduced],\n keywords=[],\n )\n\n\ndef _log1p_expander(function_expander: FunctionExpander, node: ast.Call) -> ast.AST:\n _check_num_args(node, 1)\n return ast.Call(\n func=ast.Name(id=\"log\", ctx=ast.Load()),\n args=[\n ast.BinOp(\n left=ast_utils.make_constant(1),\n op=ast.Add(),\n right=function_expander.visit(node.args[0]),\n )\n ],\n keywords=[],\n )\n\n\ndef _pow_expander(function_expander: FunctionExpander, node: ast.Call) -> ast.AST:\n _check_num_args(node, 2)\n return ast.BinOp(\n left=function_expander.visit(node.args[0]),\n op=ast.Pow(),\n right=function_expander.visit(node.args[1]),\n )\n\n\ndef _check_num_args(node: ast.Call, nargs: int) -> None:\n if len(node.args) != nargs:\n fn_name = ast_utils.extract_function_name_or_none(node)\n raise exceptions.LatexifySyntaxError(\n f\"Incorrect number of arguments for {fn_name}.\"\n f\" expected: {nargs}, but got {len(node.args)}\"\n )\n\n\n_FUNCTION_EXPANDERS: dict[str, Callable[[FunctionExpander, ast.Call], ast.AST]] = {\n \"atan2\": _atan2_expander,\n \"exp\": _exp_expander,\n \"exp2\": _exp2_expander,\n \"expm1\": _expm1_expander,\n \"hypot\": _hypot_expander,\n \"log1p\": _log1p_expander,\n \"pow\": _pow_expander,\n}\n"
},
{
"path": "src/latexify/transformers/function_expander_test.py",
"content": "\"\"\"Tests for latexify.transformers.function_expander.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\n\nfrom latexify import ast_utils, test_utils\nfrom latexify.transformers.function_expander import FunctionExpander\n\n\ndef test_preserve_keywords() -> None:\n tree = ast.Call(\n func=ast_utils.make_name(\"f\"),\n args=[ast_utils.make_name(\"x\")],\n keywords=[ast.keyword(arg=\"y\", value=ast_utils.make_constant(0))],\n )\n expected = ast.Call(\n func=ast_utils.make_name(\"f\"),\n args=[ast_utils.make_name(\"x\")],\n keywords=[ast.keyword(arg=\"y\", value=ast_utils.make_constant(0))],\n )\n transformed = FunctionExpander(set()).visit(tree)\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_exp() -> None:\n tree = ast.Call(\n func=ast_utils.make_name(\"exp\"),\n args=[ast_utils.make_name(\"x\")],\n keywords=[],\n )\n expected = ast.BinOp(\n left=ast_utils.make_name(\"e\"),\n op=ast.Pow(),\n right=ast_utils.make_name(\"x\"),\n )\n transformed = FunctionExpander({\"exp\"}).visit(tree)\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_exp_unchanged() -> None:\n tree = ast.Call(\n func=ast_utils.make_name(\"exp\"),\n args=[ast_utils.make_name(\"x\")],\n keywords=[],\n )\n expected = ast.Call(\n func=ast_utils.make_name(\"exp\"),\n args=[ast_utils.make_name(\"x\")],\n keywords=[],\n )\n transformed = FunctionExpander(set()).visit(tree)\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_exp_with_attribute() -> None:\n tree = ast.Call(\n func=ast_utils.make_attribute(ast_utils.make_name(\"math\"), \"exp\"),\n args=[ast_utils.make_name(\"x\")],\n keywords=[],\n )\n expected = ast.BinOp(\n left=ast_utils.make_name(\"e\"),\n op=ast.Pow(),\n right=ast_utils.make_name(\"x\"),\n )\n transformed2 = FunctionExpander({\"exp\"}).visit(tree)\n test_utils.assert_ast_equal(transformed2, expected)\n\n\ndef test_exp_unchanged_with_attribute() -> None:\n tree = ast.Call(\n func=ast_utils.make_attribute(ast_utils.make_name(\"math\"), \"exp\"),\n args=[ast_utils.make_name(\"x\")],\n keywords=[],\n )\n expected = ast.Call(\n func=ast_utils.make_attribute(ast_utils.make_name(\"math\"), \"exp\"),\n args=[ast_utils.make_name(\"x\")],\n keywords=[],\n )\n transformed = FunctionExpander(set()).visit(tree)\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_exp_nested1() -> None:\n tree = ast.Call(\n func=ast_utils.make_name(\"exp\"),\n args=[\n ast.Call(\n func=ast_utils.make_name(\"exp\"),\n args=[ast_utils.make_name(\"x\")],\n keywords=[],\n )\n ],\n keywords=[],\n )\n expected = ast.BinOp(\n left=ast_utils.make_name(\"e\"),\n op=ast.Pow(),\n right=ast.BinOp(\n left=ast_utils.make_name(\"e\"),\n op=ast.Pow(),\n right=ast_utils.make_name(\"x\"),\n ),\n )\n transformed = FunctionExpander({\"exp\"}).visit(tree)\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_exp_nested2() -> None:\n tree = ast.Call(\n func=ast_utils.make_name(\"f\"),\n args=[\n ast.Call(\n func=ast_utils.make_name(\"exp\"),\n args=[ast_utils.make_name(\"x\")],\n keywords=[],\n )\n ],\n keywords=[],\n )\n expected = ast.Call(\n func=ast_utils.make_name(\"f\"),\n args=[\n ast.BinOp(\n left=ast_utils.make_name(\"e\"),\n op=ast.Pow(),\n right=ast_utils.make_name(\"x\"),\n )\n ],\n keywords=[],\n )\n transformed = FunctionExpander({\"exp\"}).visit(tree)\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_atan2() -> None:\n tree = ast.Call(\n func=ast_utils.make_name(\"atan2\"),\n args=[ast_utils.make_name(\"y\"), ast_utils.make_name(\"x\")],\n keywords=[],\n )\n expected = ast.Call(\n func=ast_utils.make_name(\"atan\"),\n args=[\n ast.BinOp(\n left=ast_utils.make_name(\"y\"),\n op=ast.Div(),\n right=ast_utils.make_name(\"x\"),\n )\n ],\n keywords=[],\n )\n transformed = FunctionExpander({\"atan2\"}).visit(tree)\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_exp2() -> None:\n tree = ast.Call(\n func=ast_utils.make_name(\"exp2\"),\n args=[ast_utils.make_name(\"x\")],\n keywords=[],\n )\n expected = ast.BinOp(\n left=ast_utils.make_constant(2),\n op=ast.Pow(),\n right=ast_utils.make_name(\"x\"),\n )\n transformed = FunctionExpander({\"exp2\"}).visit(tree)\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_expm1() -> None:\n tree = ast.Call(\n func=ast_utils.make_name(\"expm1\"),\n args=[ast_utils.make_name(\"x\")],\n keywords=[],\n )\n expected = ast.BinOp(\n left=ast.Call(\n func=ast_utils.make_name(\"exp\"),\n args=[ast_utils.make_name(\"x\")],\n keywords=[],\n ),\n op=ast.Sub(),\n right=ast_utils.make_constant(1),\n )\n transformed = FunctionExpander({\"expm1\"}).visit(tree)\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_hypot() -> None:\n tree = ast.Call(\n func=ast_utils.make_name(\"hypot\"),\n args=[ast_utils.make_name(\"x\"), ast_utils.make_name(\"y\")],\n keywords=[],\n )\n expected = ast.Call(\n func=ast_utils.make_name(\"sqrt\"),\n args=[\n ast.BinOp(\n left=ast.BinOp(\n left=ast_utils.make_name(\"x\"),\n op=ast.Pow(),\n right=ast_utils.make_constant(2),\n ),\n op=ast.Add(),\n right=ast.BinOp(\n left=ast_utils.make_name(\"y\"),\n op=ast.Pow(),\n right=ast_utils.make_constant(2),\n ),\n )\n ],\n keywords=[],\n )\n transformed = FunctionExpander({\"hypot\"}).visit(tree)\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_hypot_no_args() -> None:\n tree = ast.Call(func=ast_utils.make_name(\"hypot\"), args=[], keywords=[])\n expected = ast_utils.make_constant(0)\n transformed = FunctionExpander({\"hypot\"}).visit(tree)\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_log1p() -> None:\n tree = ast.Call(\n func=ast_utils.make_name(\"log1p\"),\n args=[ast_utils.make_name(\"x\")],\n keywords=[],\n )\n expected = ast.Call(\n func=ast_utils.make_name(\"log\"),\n args=[\n ast.BinOp(\n left=ast_utils.make_constant(1),\n op=ast.Add(),\n right=ast_utils.make_name(\"x\"),\n )\n ],\n keywords=[],\n )\n transformed = FunctionExpander({\"log1p\"}).visit(tree)\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_pow() -> None:\n tree = ast.Call(\n func=ast_utils.make_name(\"pow\"),\n args=[ast_utils.make_name(\"x\"), ast_utils.make_name(\"y\")],\n keywords=[],\n )\n expected = ast.BinOp(\n left=ast_utils.make_name(\"x\"),\n op=ast.Pow(),\n right=ast_utils.make_name(\"y\"),\n )\n transformed = FunctionExpander({\"pow\"}).visit(tree)\n test_utils.assert_ast_equal(transformed, expected)\n"
},
{
"path": "src/latexify/transformers/identifier_replacer.py",
"content": "\"\"\"Transformer to replace user symbols.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\nimport keyword\nfrom typing import cast\n\nfrom latexify import ast_utils\n\n\nclass IdentifierReplacer(ast.NodeTransformer):\n \"\"\"NodeTransformer to replace identifier names.\n\n This class defines a rule to replace identifiers in AST with specified names.\n\n Example:\n def foo(bar):\n return baz\n\n IdentifierReplacer({\"foo\": \"x\", \"bar\": \"y\", \"baz\": \"z\"}) will modify the AST of\n the function above to below:\n\n def x(y):\n return z\n \"\"\"\n\n def __init__(self, mapping: dict[str, str]):\n \"\"\"Initializer.\n\n Args:\n mapping: User defined mapping of names. Keys are the original names of the\n identifiers, and corresponding values are the replacements.\n Both keys and values have to represent valid Python identifiers:\n ^[A-Za-z_][A-Za-z0-9_]*$\n \"\"\"\n self._mapping = mapping\n\n for k, v in self._mapping.items():\n if not str.isidentifier(k) or keyword.iskeyword(k):\n raise ValueError(f\"'{k}' is not an identifier name.\")\n if not str.isidentifier(v) or keyword.iskeyword(v):\n raise ValueError(f\"'{v}' is not an identifier name.\")\n\n def _replace_args(self, args: list[ast.arg]) -> list[ast.arg]:\n \"\"\"Helper function to replace arg names.\"\"\"\n return [ast.arg(arg=self._mapping.get(a.arg, a.arg)) for a in args]\n\n def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:\n \"\"\"Visit a FunctionDef node.\"\"\"\n visited = cast(ast.FunctionDef, super().generic_visit(node))\n\n args = ast.arguments(\n posonlyargs=self._replace_args(visited.args.posonlyargs),\n args=self._replace_args(visited.args.args),\n vararg=visited.args.vararg,\n kwonlyargs=self._replace_args(visited.args.kwonlyargs),\n kw_defaults=visited.args.kw_defaults,\n kwarg=visited.args.kwarg,\n defaults=visited.args.defaults,\n )\n type_params = getattr(visited, \"type_params\", [])\n return ast_utils.create_function_def(\n name=self._mapping.get(visited.name, visited.name),\n args=args,\n body=visited.body,\n decorator_list=visited.decorator_list,\n returns=visited.returns,\n type_params=type_params,\n )\n\n def visit_Name(self, node: ast.Name) -> ast.Name:\n \"\"\"Visit a Name node.\"\"\"\n return ast.Name(\n id=self._mapping.get(node.id, node.id),\n ctx=node.ctx,\n )\n"
},
{
"path": "src/latexify/transformers/identifier_replacer_test.py",
"content": "\"\"\"Tests for latexify.transformer.identifier_replacer.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\n\nimport pytest\n\nfrom latexify import ast_utils, test_utils\nfrom latexify.transformers.identifier_replacer import IdentifierReplacer\n\n\ndef test_invalid_mapping() -> None:\n with pytest.raises(ValueError, match=r\"'123' is not an identifier name.\"):\n IdentifierReplacer({\"123\": \"foo\"})\n with pytest.raises(ValueError, match=r\"'456' is not an identifier name.\"):\n IdentifierReplacer({\"foo\": \"456\"})\n with pytest.raises(ValueError, match=r\"'def' is not an identifier name.\"):\n IdentifierReplacer({\"foo\": \"def\"})\n\n\ndef test_name_replaced() -> None:\n source = ast.Name(id=\"foo\", ctx=ast.Load())\n expected = ast.Name(id=\"bar\", ctx=ast.Load())\n transformed = IdentifierReplacer({\"foo\": \"bar\"}).visit(source)\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_name_not_replaced() -> None:\n source = ast.Name(id=\"foo\", ctx=ast.Load())\n expected = ast.Name(id=\"foo\", ctx=ast.Load())\n transformed = IdentifierReplacer({\"fo\": \"bar\"}).visit(source)\n test_utils.assert_ast_equal(transformed, expected)\n transformed = IdentifierReplacer({\"fooo\": \"bar\"}).visit(source)\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_functiondef_with_posonlyargs() -> None:\n # Subtree of:\n # @d\n # def f(x=a, /, y=b, *, z=c):\n # pass\n source = ast_utils.create_function_def(\n name=\"f\",\n args=ast.arguments(\n posonlyargs=[ast.arg(arg=\"x\")],\n args=[ast.arg(arg=\"y\")],\n kwonlyargs=[ast.arg(arg=\"z\")],\n kw_defaults=[ast.Name(id=\"c\", ctx=ast.Load())],\n defaults=[\n ast.Name(id=\"a\", ctx=ast.Load()),\n ast.Name(id=\"b\", ctx=ast.Load()),\n ],\n ),\n body=[ast.Pass()],\n decorator_list=[ast.Name(id=\"d\", ctx=ast.Load())],\n returns=None,\n type_comment=None,\n type_params=[],\n )\n\n expected = ast_utils.create_function_def(\n name=\"F\",\n args=ast.arguments(\n posonlyargs=[ast.arg(arg=\"X\")],\n args=[ast.arg(arg=\"Y\")],\n kwonlyargs=[ast.arg(arg=\"Z\")],\n kw_defaults=[ast.Name(id=\"C\", ctx=ast.Load())],\n defaults=[\n ast.Name(id=\"A\", ctx=ast.Load()),\n ast.Name(id=\"B\", ctx=ast.Load()),\n ],\n ),\n body=[ast.Pass()],\n decorator_list=[ast.Name(id=\"D\", ctx=ast.Load())],\n returns=None,\n type_comment=None,\n type_params=[],\n )\n\n mapping = {x: x.upper() for x in \"abcdfxyz\"}\n transformed = IdentifierReplacer(mapping).visit(source)\n test_utils.assert_ast_equal(transformed, expected)\n\n\ndef test_expr() -> None:\n # Subtree of:\n # (x + y) * z\n source = ast.BinOp(\n left=ast.BinOp(\n left=ast.Name(id=\"x\", ctx=ast.Load()),\n op=ast.Add(),\n right=ast.Name(id=\"y\", ctx=ast.Load()),\n ),\n op=ast.Mult(),\n right=ast.Name(id=\"z\", ctx=ast.Load()),\n )\n\n expected = ast.BinOp(\n left=ast.BinOp(\n left=ast.Name(id=\"X\", ctx=ast.Load()),\n op=ast.Add(),\n right=ast.Name(id=\"Y\", ctx=ast.Load()),\n ),\n op=ast.Mult(),\n right=ast.Name(id=\"Z\", ctx=ast.Load()),\n )\n\n mapping = {x: x.upper() for x in \"xyz\"}\n transformed = IdentifierReplacer(mapping).visit(source)\n test_utils.assert_ast_equal(transformed, expected)\n"
},
{
"path": "src/latexify/transformers/prefix_trimmer.py",
"content": "\"\"\"NodeTransformer to trim unnecessary prefixes.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\nimport re\n\nfrom latexify import ast_utils\n\n_PREFIX_PATTERN = re.compile(r\"^[A-Za-z_][A-Za-z0-9_]*(\\.[A-Za-z_][A-Za-z0-9_]*)*$\")\n\n\nclass PrefixTrimmer(ast.NodeTransformer):\n \"\"\"NodeTransformer to trim unnecessary prefixes.\n\n This class investigates all Attribute subtrees, and replace them if the prefix of\n the attribute matches the given set of prefixes.\n Prefix is searched in the manner of leftmost longest matching.\n\n Example:\n def f(x):\n return math.sqrt(x)\n\n PrefixTrimmer({\"math\"}) will modify the AST of the function above to below:\n\n def f(x):\n return sqrt(x)\n \"\"\"\n\n _prefixes: list[tuple[str, ...]]\n\n def __init__(self, prefixes: set[str]) -> None:\n \"\"\"Initializer.\n\n Args:\n prefixes: Set of prefixes to be trimmed. Nested prefix is allowed too.\n Each value must follow one of the following formats:\n - A Python identifier, e.g., \"math\"\n - Python identifiers joined by periods, e.g., \"numpy.random\"\n \"\"\"\n for p in prefixes:\n if not _PREFIX_PATTERN.match(p):\n raise ValueError(f\"Invalid prefix: {p}\")\n\n self._prefixes = [tuple(p.split(\".\")) for p in prefixes]\n\n def _get_prefix(self, node: ast.expr) -> tuple[str, ...] | None:\n \"\"\"Helper to obtain nested prefix.\n\n Args:\n node: Node to investigate.\n\n Returns:\n The prefix tuple, or None if the node has unsupported syntax.\n \"\"\"\n if isinstance(node, ast.Name):\n return (node.id,)\n\n if isinstance(node, ast.Attribute):\n parent = self._get_prefix(node.value)\n return parent + (node.attr,) if parent is not None else None\n\n return None\n\n def _make_attribute(self, prefix: tuple[str, ...], name: str) -> ast.expr:\n \"\"\"Helper to generate a new Attribute or Name node.\n\n Args:\n prefix: List of prefixes.\n name: Attribute name.\n\n Returns:\n Name node if prefix == (), (possibly nested) Attribute node otherwise.\n \"\"\"\n if not prefix:\n return ast_utils.make_name(name)\n\n parent = self._make_attribute(prefix[:-1], prefix[-1])\n return ast_utils.make_attribute(parent, name)\n\n def visit_Attribute(self, node: ast.Attribute) -> ast.expr:\n \"\"\"Visit an Attribute node.\"\"\"\n prefix = self._get_prefix(node.value)\n if prefix is None:\n return node\n\n # Performs leftmost longest match.\n # NOTE(odashi):\n # This implementation is very naive, but would work efficiently as long as the\n # number of patterns is small.\n matched_length = 0\n\n for p in self._prefixes:\n length = min(len(p), len(prefix))\n if prefix[:length] == p and length > matched_length:\n matched_length = length\n\n return self._make_attribute(prefix[matched_length:], node.attr)\n"
},
{
"path": "src/latexify/transformers/prefix_trimmer_test.py",
"content": "\"\"\"Tests for latexify.transformers.prefix_trimmer.\"\"\"\n\nfrom __future__ import annotations\n\nimport ast\n\nimport pytest\n\nfrom latexify import ast_utils, test_utils\nfrom latexify.transformers import prefix_trimmer\n\n# For convenience\nmake_name = ast_utils.make_name\nmake_attr = ast_utils.make_attribute\nPrefixTrimmer = prefix_trimmer.PrefixTrimmer\n\n\n@pytest.mark.parametrize(\n \"prefix\", [\".x\", \"x.\", \"1\", \"1x\", \"x.1\", \"x.1x\", \"x.x.1\", \"x.x.1x\" \"x..x\", \"x.x..x\"]\n)\ndef test_invalid_prefix(prefix: str) -> None:\n with pytest.raises(ValueError, match=rf\"^Invalid prefix: {prefix}$\"):\n PrefixTrimmer({prefix})\n\n\n@pytest.mark.parametrize(\n \"prefixes,expected\",\n [\n (set(), make_name(\"foo\")),\n ({\"foo\"}, make_name(\"foo\")),\n ({\"bar\"}, make_name(\"foo\")),\n ({\"foo.bar\"}, make_name(\"foo\")),\n ({\"foo\", \"bar\"}, make_name(\"foo\")),\n ({\"foo\", \"foo.bar\"}, make_name(\"foo\")),\n ],\n)\ndef test_name(prefixes: set[str], expected: ast.expr) -> None:\n source = make_name(\"foo\")\n transformed = PrefixTrimmer(prefixes).visit(source)\n test_utils.assert_ast_equal(transformed, expected)\n\n\n@pytest.mark.parametrize(\n \"prefixes,expected\",\n [\n (set(), make_attr(make_name(\"foo\"), \"bar\")),\n ({\"fo\"}, make_attr(make_name(\"foo\"), \"bar\")),\n ({\"foo\"}, make_name(\"bar\")),\n ({\"bar\"}, make_attr(make_name(\"foo\"), \"bar\")),\n ({\"baz\"}, make_attr(make_name(\"foo\"), \"bar\")),\n ({\"foo.bar\"}, make_attr(make_name(\"foo\"), \"bar\")),\n ({\"foo\", \"bar\"}, make_name(\"bar\")),\n ({\"foo\", \"foo.bar\"}, make_name(\"bar\")),\n ],\n)\ndef test_attr_1(prefixes: set[str], expected: ast.expr) -> None:\n source = make_attr(make_name(\"foo\"), \"bar\")\n transformed = PrefixTrimmer(prefixes).visit(source)\n test_utils.assert_ast_equal(transformed, expected)\n\n\n@pytest.mark.parametrize(\n \"prefixes,expected\",\n [\n (set(), make_attr(make_attr(make_name(\"foo\"), \"bar\"), \"baz\")),\n ({\"fo\"}, make_attr(make_attr(make_name(\"foo\"), \"bar\"), \"baz\")),\n ({\"foo\"}, make_attr(make_name(\"bar\"), \"baz\")),\n ({\"bar\"}, make_attr(make_attr(make_name(\"foo\"), \"bar\"), \"baz\")),\n ({\"baz\"}, make_attr(make_attr(make_name(\"foo\"), \"bar\"), \"baz\")),\n ({\"foo.ba\"}, make_attr(make_attr(make_name(\"foo\"), \"bar\"), \"baz\")),\n ({\"foo.bar\"}, make_name(\"baz\")),\n ({\"foo.bar.baz\"}, make_attr(make_attr(make_name(\"foo\"), \"bar\"), \"baz\")),\n ({\"foo\", \"bar\"}, make_attr(make_name(\"bar\"), \"baz\")),\n ({\"foo\", \"foo.bar\"}, make_name(\"baz\")),\n ],\n)\ndef test_attr_2(prefixes: set[str], expected: ast.expr) -> None:\n source = make_attr(make_attr(make_name(\"foo\"), \"bar\"), \"baz\")\n transformed = PrefixTrimmer(prefixes).visit(source)\n test_utils.assert_ast_equal(transformed, expected)\n"
}
]