From 97127fad4d02de2bcadf611e75e066544bd99a98 Mon Sep 17 00:00:00 2001 From: nstarman Date: Wed, 10 Jul 2024 17:38:30 -0400 Subject: [PATCH] initial commit Signed-off-by: nstarman --- .copier-answers.yml | 12 +++ .git_archival.txt | 4 + .gitattributes | 1 + .github/CONTRIBUTING.md | 101 ++++++++++++++++++ .github/dependabot.yml | 11 ++ .github/matchers/pylint.json | 32 ++++++ .github/workflows/cd.yml | 54 ++++++++++ .github/workflows/ci.yml | 73 +++++++++++++ .gitignore | 158 ++++++++++++++++++++++++++++ .pre-commit-config.yaml | 87 ++++++++++++++++ .readthedocs.yaml | 18 ++++ LICENSE | 19 ++++ README.md | 27 +++++ docs/conf.py | 45 ++++++++ docs/index.md | 17 +++ noxfile.py | 119 +++++++++++++++++++++ pyproject.toml | 160 +++++++++++++++++++++++++++++ src/immutable_map_jax/__init__.py | 10 ++ src/immutable_map_jax/_core.py | 140 +++++++++++++++++++++++++ src/immutable_map_jax/_version.pyi | 4 + src/immutable_map_jax/py.typed | 0 tests/test_package.py | 9 ++ 22 files changed, 1101 insertions(+) create mode 100644 .copier-answers.yml create mode 100644 .git_archival.txt create mode 100644 .gitattributes create mode 100644 .github/CONTRIBUTING.md create mode 100644 .github/dependabot.yml create mode 100644 .github/matchers/pylint.json create mode 100644 .github/workflows/cd.yml create mode 100644 .github/workflows/ci.yml create mode 100644 .gitignore create mode 100644 .pre-commit-config.yaml create mode 100644 .readthedocs.yaml create mode 100644 LICENSE create mode 100644 README.md create mode 100644 docs/conf.py create mode 100644 docs/index.md create mode 100644 noxfile.py create mode 100644 pyproject.toml create mode 100644 src/immutable_map_jax/__init__.py create mode 100644 src/immutable_map_jax/_core.py create mode 100644 src/immutable_map_jax/_version.pyi create mode 100644 src/immutable_map_jax/py.typed create mode 100644 tests/test_package.py diff --git a/.copier-answers.yml b/.copier-answers.yml new file mode 100644 index 0000000..b0da338 --- /dev/null +++ b/.copier-answers.yml @@ -0,0 +1,12 @@ +# Changes here will be overwritten by Copier; NEVER EDIT MANUALLY +_commit: 2024.04.23 +_src_path: gh:scientific-python/cookie +backend: hatch +email: nstarman@users.noreply.github.com +full_name: Galactic Dynamics Maintainers +license: MIT +org: GalacticDynamics +project_name: immutable_map_jax +project_short_description: Immutable Map, compatible with JAX & Equinox +url: https://github.com/GalacticDynamics/immutable_map_jax +vcs: true diff --git a/.git_archival.txt b/.git_archival.txt new file mode 100644 index 0000000..8fb235d --- /dev/null +++ b/.git_archival.txt @@ -0,0 +1,4 @@ +node: $Format:%H$ +node-date: $Format:%cI$ +describe-name: $Format:%(describe:tags=true,match=*[0-9]*)$ +ref-names: $Format:%D$ diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..00a7b00 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +.git_archival.txt export-subst diff --git a/.github/CONTRIBUTING.md b/.github/CONTRIBUTING.md new file mode 100644 index 0000000..cac7282 --- /dev/null +++ b/.github/CONTRIBUTING.md @@ -0,0 +1,101 @@ +See the [Scientific Python Developer Guide][spc-dev-intro] for a detailed +description of best practices for developing scientific packages. + +[spc-dev-intro]: https://learn.scientific-python.org/development/ + +# Quick development + +The fastest way to start with development is to use nox. If you don't have nox, +you can use `pipx run nox` to run it without installing, or `pipx install nox`. +If you don't have pipx (pip for applications), then you can install with +`pip install pipx` (the only case were installing an application with regular +pip is reasonable). If you use macOS, then pipx and nox are both in brew, use +`brew install pipx nox`. + +To use, run `nox`. This will lint and test using every installed version of +Python on your system, skipping ones that are not installed. You can also run +specific jobs: + +```console +$ nox -s lint # Lint only +$ nox -s tests # Python tests +$ nox -s docs -- --serve # Build and serve the docs +$ nox -s build # Make an SDist and wheel +``` + +Nox handles everything for you, including setting up an temporary virtual +environment for each run. + +# Setting up a development environment manually + +You can set up a development environment by running: + +```bash +python3 -m venv .venv +source ./.venv/bin/activate +pip install -v -e .[dev] +``` + +If you have the +[Python Launcher for Unix](https://github.com/brettcannon/python-launcher), you +can instead do: + +```bash +py -m venv .venv +py -m install -v -e .[dev] +``` + +# Post setup + +You should prepare pre-commit, which will help you by checking that commits pass +required checks: + +```bash +pip install pre-commit # or brew install pre-commit on macOS +pre-commit install # Will install a pre-commit hook into the git repo +``` + +You can also/alternatively run `pre-commit run` (changes only) or +`pre-commit run --all-files` to check even without installing the hook. + +# Testing + +Use pytest to run the unit checks: + +```bash +pytest +``` + +# Coverage + +Use pytest-cov to generate coverage reports: + +```bash +pytest --cov=immutable_map_jax +``` + +# Building docs + +You can build the docs using: + +```bash +nox -s docs +``` + +You can see a preview with: + +```bash +nox -s docs -- --serve +``` + +# Pre-commit + +This project uses pre-commit for all style checking. While you can run it with +nox, this is such an important tool that it deserves to be installed on its own. +Install pre-commit and run: + +```bash +pre-commit run -a +``` + +to check all files. diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..6c4b369 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,11 @@ +version: 2 +updates: + # Maintain dependencies for GitHub Actions + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "weekly" + groups: + actions: + patterns: + - "*" diff --git a/.github/matchers/pylint.json b/.github/matchers/pylint.json new file mode 100644 index 0000000..e3a6bd1 --- /dev/null +++ b/.github/matchers/pylint.json @@ -0,0 +1,32 @@ +{ + "problemMatcher": [ + { + "severity": "warning", + "pattern": [ + { + "regexp": "^([^:]+):(\\d+):(\\d+): ([A-DF-Z]\\d+): \\033\\[[\\d;]+m([^\\033]+).*$", + "file": 1, + "line": 2, + "column": 3, + "code": 4, + "message": 5 + } + ], + "owner": "pylint-warning" + }, + { + "severity": "error", + "pattern": [ + { + "regexp": "^([^:]+):(\\d+):(\\d+): (E\\d+): \\033\\[[\\d;]+m([^\\033]+).*$", + "file": 1, + "line": 2, + "column": 3, + "code": 4, + "message": 5 + } + ], + "owner": "pylint-error" + } + ] +} diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml new file mode 100644 index 0000000..ead459d --- /dev/null +++ b/.github/workflows/cd.yml @@ -0,0 +1,54 @@ +name: CD + +on: + workflow_dispatch: + pull_request: + push: + branches: + - main + release: + types: + - published + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +env: + # Many color libraries just need this to be set to any value, but at least + # one distinguishes color depth, where "3" -> "256-bit color". + FORCE_COLOR: 3 + +jobs: + dist: + name: Distribution build + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - uses: hynek/build-and-inspect-python-package@v2 + + publish: + needs: [dist] + name: Publish to PyPI + environment: pypi + permissions: + id-token: write + runs-on: ubuntu-latest + if: github.event_name == 'release' && github.event.action == 'published' + + steps: + - uses: actions/download-artifact@v4 + with: + name: Packages + path: dist + + - uses: pypa/gh-action-pypi-publish@release/v1 + if: github.event_name == 'release' && github.event.action == 'published' + with: + # Remember to tell (test-)pypi about this repo before publishing + # Remove this line to publish to PyPI + repository-url: https://test.pypi.org/legacy/ diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..7898bd0 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,73 @@ +name: CI + +on: + workflow_dispatch: + pull_request: + push: + branches: + - main + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +env: + # Many color libraries just need this to be set to any value, but at least + # one distinguishes color depth, where "3" -> "256-bit color". + FORCE_COLOR: 3 + +jobs: + pre-commit: + name: Format + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - uses: actions/setup-python@v5 + with: + python-version: "3.x" + - uses: pre-commit/action@v3.0.1 + with: + extra_args: --hook-stage manual --all-files + - name: Run PyLint + run: | + echo "::add-matcher::$GITHUB_WORKSPACE/.github/matchers/pylint.json" + pipx run nox -s pylint + + checks: + name: Check Python ${{ matrix.python-version }} on ${{ matrix.runs-on }} + runs-on: ${{ matrix.runs-on }} + needs: [pre-commit] + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.12"] + runs-on: [ubuntu-latest, macos-latest, windows-latest] + + include: + - python-version: pypy-3.10 + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + allow-prereleases: true + + - name: Install package + run: python -m pip install .[test] + + - name: Test package + run: >- + python -m pytest -ra --cov --cov-report=xml --cov-report=term + --durations=20 + + - name: Upload coverage report + uses: codecov/codecov-action@v4.3.0 + with: + token: ${{ secrets.CODECOV_TOKEN }} diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..25cf9a4 --- /dev/null +++ b/.gitignore @@ -0,0 +1,158 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# setuptools_scm +src/*/_version.py + + +# ruff +.ruff_cache/ + +# OS specific stuff +.DS_Store +.DS_Store? +._* +.Spotlight-V100 +.Trashes +ehthumbs.db +Thumbs.db + +# Common editor files +*~ +*.swp diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..9fc6fdc --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,87 @@ +ci: + autoupdate_commit_msg: "chore: update pre-commit hooks" + autofix_commit_msg: "style: pre-commit fixes" + +repos: + - repo: https://github.com/adamchainz/blacken-docs + rev: "1.16.0" + hooks: + - id: blacken-docs + additional_dependencies: [black==24.*] + + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: "v4.6.0" + hooks: + - id: check-added-large-files + - id: check-case-conflict + - id: check-merge-conflict + - id: check-symlinks + - id: check-yaml + - id: debug-statements + - id: end-of-file-fixer + - id: mixed-line-ending + - id: name-tests-test + args: ["--pytest-test-first"] + - id: requirements-txt-fixer + - id: trailing-whitespace + + - repo: https://github.com/pre-commit/pygrep-hooks + rev: "v1.10.0" + hooks: + - id: rst-backticks + - id: rst-directive-colons + - id: rst-inline-touching-normal + + - repo: https://github.com/pre-commit/mirrors-prettier + rev: "v3.1.0" + hooks: + - id: prettier + types_or: [yaml, markdown, html, css, scss, javascript, json] + args: [--prose-wrap=always] + + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: "v0.4.1" + hooks: + - id: ruff + args: ["--fix", "--show-fixes"] + - id: ruff-format + + - repo: https://github.com/pre-commit/mirrors-mypy + rev: "v1.9.0" + hooks: + - id: mypy + files: src|tests + args: [] + additional_dependencies: + - pytest + + - repo: https://github.com/codespell-project/codespell + rev: "v2.2.6" + hooks: + - id: codespell + + - repo: https://github.com/shellcheck-py/shellcheck-py + rev: "v0.10.0.1" + hooks: + - id: shellcheck + + - repo: local + hooks: + - id: disallow-caps + name: Disallow improper capitalization + language: pygrep + entry: PyBind|Numpy|Cmake|CCache|Github|PyTest + exclude: .pre-commit-config.yaml + + - repo: https://github.com/abravalheri/validate-pyproject + rev: "v0.16" + hooks: + - id: validate-pyproject + additional_dependencies: ["validate-pyproject-schema-store[all]"] + + - repo: https://github.com/python-jsonschema/check-jsonschema + rev: "0.28.2" + hooks: + - id: check-dependabot + - id: check-github-workflows + - id: check-readthedocs diff --git a/.readthedocs.yaml b/.readthedocs.yaml new file mode 100644 index 0000000..7e49657 --- /dev/null +++ b/.readthedocs.yaml @@ -0,0 +1,18 @@ +# Read the Docs configuration file +# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details + +version: 2 + +build: + os: ubuntu-22.04 + tools: + python: "3.11" +sphinx: + configuration: docs/conf.py + +python: + install: + - method: pip + path: . + extra_requirements: + - docs diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..8099acf --- /dev/null +++ b/LICENSE @@ -0,0 +1,19 @@ +Copyright 2024 Galactic Dynamics Maintainers + +Permission is hereby granted, free of charge, to any person obtaining a copy of +this software and associated documentation files (the "Software"), to deal in +the Software without restriction, including without limitation the rights to +use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies +of the Software, and to permit persons to whom the Software is furnished to do +so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..9be62b2 --- /dev/null +++ b/README.md @@ -0,0 +1,27 @@ +# immutable_map_jax + +[![Actions Status][actions-badge]][actions-link] +[![Documentation Status][rtd-badge]][rtd-link] + +[![PyPI version][pypi-version]][pypi-link] +[![Conda-Forge][conda-badge]][conda-link] +[![PyPI platforms][pypi-platforms]][pypi-link] + +[![GitHub Discussion][github-discussions-badge]][github-discussions-link] + + + + +[actions-badge]: https://github.com/GalacticDynamics/immutable_map_jax/workflows/CI/badge.svg +[actions-link]: https://github.com/GalacticDynamics/immutable_map_jax/actions +[conda-badge]: https://img.shields.io/conda/vn/conda-forge/immutable_map_jax +[conda-link]: https://github.com/conda-forge/immutable_map_jax-feedstock +[github-discussions-badge]: https://img.shields.io/static/v1?label=Discussions&message=Ask&color=blue&logo=github +[github-discussions-link]: https://github.com/GalacticDynamics/immutable_map_jax/discussions +[pypi-link]: https://pypi.org/project/immutable_map_jax/ +[pypi-platforms]: https://img.shields.io/pypi/pyversions/immutable_map_jax +[pypi-version]: https://img.shields.io/pypi/v/immutable_map_jax +[rtd-badge]: https://readthedocs.org/projects/immutable_map_jax/badge/?version=latest +[rtd-link]: https://immutable_map_jax.readthedocs.io/en/latest/?badge=latest + + diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 index 0000000..27e8d36 --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +import importlib.metadata + +project = "immutable_map_jax" +copyright = "2024, Galactic Dynamics Maintainers" +author = "Galactic Dynamics Maintainers" +version = release = importlib.metadata.version("immutable_map_jax") + +extensions = [ + "myst_parser", + "sphinx.ext.autodoc", + "sphinx.ext.intersphinx", + "sphinx.ext.mathjax", + "sphinx.ext.napoleon", + "sphinx_autodoc_typehints", + "sphinx_copybutton", +] + +source_suffix = [".rst", ".md"] +exclude_patterns = [ + "_build", + "**.ipynb_checkpoints", + "Thumbs.db", + ".DS_Store", + ".env", + ".venv", +] + +html_theme = "furo" + +myst_enable_extensions = [ + "colon_fence", +] + +intersphinx_mapping = { + "python": ("https://docs.python.org/3", None), +} + +nitpick_ignore = [ + ("py:class", "_io.StringIO"), + ("py:class", "_io.BytesIO"), +] + +always_document_param_types = True diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 0000000..37151fb --- /dev/null +++ b/docs/index.md @@ -0,0 +1,17 @@ +# immutable_map_jax + +```{toctree} +:maxdepth: 2 +:hidden: + +``` + +```{include} ../README.md +:start-after: +``` + +## Indices and tables + +- {ref}`genindex` +- {ref}`modindex` +- {ref}`search` diff --git a/noxfile.py b/noxfile.py new file mode 100644 index 0000000..b3e7c87 --- /dev/null +++ b/noxfile.py @@ -0,0 +1,119 @@ +from __future__ import annotations + +import argparse +import shutil +from pathlib import Path + +import nox + +DIR = Path(__file__).parent.resolve() + +nox.needs_version = ">=2024.3.2" +nox.options.sessions = ["lint", "pylint", "tests"] +nox.options.default_venv_backend = "uv|virtualenv" + + +@nox.session +def lint(session: nox.Session) -> None: + """ + Run the linter. + """ + session.install("pre-commit") + session.run( + "pre-commit", "run", "--all-files", "--show-diff-on-failure", *session.posargs + ) + + +@nox.session +def pylint(session: nox.Session) -> None: + """ + Run PyLint. + """ + # This needs to be installed into the package environment, and is slower + # than a pre-commit check + session.install(".", "pylint") + session.run("pylint", "immutable_map_jax", *session.posargs) + + +@nox.session +def tests(session: nox.Session) -> None: + """ + Run the unit and regular tests. + """ + session.install(".[test]") + session.run("pytest", *session.posargs) + + +@nox.session(reuse_venv=True) +def docs(session: nox.Session) -> None: + """ + Build the docs. Pass "--serve" to serve. Pass "-b linkcheck" to check links. + """ + + parser = argparse.ArgumentParser() + parser.add_argument("--serve", action="store_true", help="Serve after building") + parser.add_argument( + "-b", dest="builder", default="html", help="Build target (default: html)" + ) + args, posargs = parser.parse_known_args(session.posargs) + + if args.builder != "html" and args.serve: + session.error("Must not specify non-HTML builder with --serve") + + extra_installs = ["sphinx-autobuild"] if args.serve else [] + + session.install("-e.[docs]", *extra_installs) + session.chdir("docs") + + if args.builder == "linkcheck": + session.run( + "sphinx-build", "-b", "linkcheck", ".", "_build/linkcheck", *posargs + ) + return + + shared_args = ( + "-n", # nitpicky mode + "-T", # full tracebacks + f"-b={args.builder}", + ".", + f"_build/{args.builder}", + *posargs, + ) + + if args.serve: + session.run("sphinx-autobuild", *shared_args) + else: + session.run("sphinx-build", "--keep-going", *shared_args) + + +@nox.session +def build_api_docs(session: nox.Session) -> None: + """ + Build (regenerate) API docs. + """ + + session.install("sphinx") + session.chdir("docs") + session.run( + "sphinx-apidoc", + "-o", + "api/", + "--module-first", + "--no-toc", + "--force", + "../src/immutable_map_jax", + ) + + +@nox.session +def build(session: nox.Session) -> None: + """ + Build an SDist and wheel. + """ + + build_path = DIR.joinpath("build") + if build_path.exists(): + shutil.rmtree(build_path) + + session.install("build") + session.run("python", "-m", "build") diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..7ca2fbc --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,160 @@ +[build-system] +requires = ["hatchling", "hatch-vcs"] +build-backend = "hatchling.build" + + +[project] +name = "immutable_map_jax" +authors = [ + { name = "Galactic Dynamics Maintainers", email = "nstarman@users.noreply.github.com" }, +] +description = "Immutable Map, compatible with JAX & Equinox" +readme = "README.md" +license.file = "LICENSE" +requires-python = ">=3.10" +classifiers = [ + "Development Status :: 1 - Planning", + "Intended Audience :: Science/Research", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Topic :: Scientific/Engineering", + "Typing :: Typed", +] +dynamic = ["version"] +dependencies = [ + "jax", +] + +[project.optional-dependencies] +test = [ + "pytest >=6", + "pytest-cov >=3", +] +dev = [ + "pytest >=6", + "pytest-cov >=3", +] +docs = [ + "sphinx>=7.0", + "myst_parser>=0.13", + "sphinx_copybutton", + "sphinx_autodoc_typehints", + "furo>=2023.08.17", +] + +[project.urls] +Homepage = "https://github.com/GalacticDynamics/immutable_map_jax" +"Bug Tracker" = "https://github.com/GalacticDynamics/immutable_map_jax/issues" +Discussions = "https://github.com/GalacticDynamics/immutable_map_jax/discussions" +Changelog = "https://github.com/GalacticDynamics/immutable_map_jax/releases" + + +[tool.hatch] +version.source = "vcs" +build.hooks.vcs.version-file = "src/immutable_map_jax/_version.py" + +[tool.hatch.envs.default] +features = ["test"] +scripts.test = "pytest {args}" + + +[tool.pytest.ini_options] +minversion = "6.0" +addopts = ["-ra", "--showlocals", "--strict-markers", "--strict-config"] +xfail_strict = true +filterwarnings = [ + "error", +] +log_cli_level = "INFO" +testpaths = [ + "tests", +] + + +[tool.coverage] +run.source = ["immutable_map_jax"] +report.exclude_also = [ + '\.\.\.', + 'if typing.TYPE_CHECKING:', +] + +[tool.mypy] +files = ["src"] +python_version = "3.10" +warn_unused_configs = true +strict = true +enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"] +warn_unreachable = true +disallow_untyped_defs = false +disallow_incomplete_defs = false + +[[tool.mypy.overrides]] +module = "immutable_map_jax.*" +disallow_untyped_defs = true +disallow_incomplete_defs = true + +[[tool.mypy.overrides]] +module = "jax.*" +ignore_missing_imports = true + + +[tool.ruff] +src = ["src"] + +[tool.ruff.lint] +extend-select = [ + "B", # flake8-bugbear + "I", # isort + "ARG", # flake8-unused-arguments + "C4", # flake8-comprehensions + "EM", # flake8-errmsg + "ICN", # flake8-import-conventions + "G", # flake8-logging-format + "PGH", # pygrep-hooks + "PIE", # flake8-pie + "PL", # pylint + "PT", # flake8-pytest-style + "PTH", # flake8-use-pathlib + "RET", # flake8-return + "RUF", # Ruff-specific + "SIM", # flake8-simplify + "T20", # flake8-print + "UP", # pyupgrade + "YTT", # flake8-2020 + "EXE", # flake8-executable + "NPY", # NumPy specific rules + "PD", # pandas-vet +] +ignore = [ + "PLR09", # Too many <...> + "PLR2004", # Magic value used in comparison + "ISC001", # Conflicts with formatter +] + +[tool.ruff.lint.per-file-ignores] +"tests/**" = ["T20"] +"noxfile.py" = ["T20"] + + +[tool.pylint] +py-version = "3.10" +ignore-paths = [".*/_version.py"] +reports.output-format = "colorized" +similarities.ignore-imports = "yes" +messages_control.disable = [ + "design", + "fixme", + "line-too-long", + "missing-module-docstring", + "missing-function-docstring", + "wrong-import-position", +] diff --git a/src/immutable_map_jax/__init__.py b/src/immutable_map_jax/__init__.py new file mode 100644 index 0000000..2218b27 --- /dev/null +++ b/src/immutable_map_jax/__init__.py @@ -0,0 +1,10 @@ +""" +Copyright (c) 2024 Galactic Dynamics Maintainers. All rights reserved. + +immutable_map_jax: Immutable Map, compatible with JAX & Equinox +""" + +__all__ = ["__version__", "ImmutableMap"] + +from ._core import ImmutableMap +from ._version import version as __version__ diff --git a/src/immutable_map_jax/_core.py b/src/immutable_map_jax/_core.py new file mode 100644 index 0000000..daebd4c --- /dev/null +++ b/src/immutable_map_jax/_core.py @@ -0,0 +1,140 @@ +""" +Copyright (c) 2024 Galactic Dynamics Maintainers. All rights reserved. + +immutable_map_jax: Immutable Map, compatible with JAX & Equinox +""" + +__all__: list[str] = [] + +from collections.abc import ItemsView, Iterable, Iterator, KeysView, Mapping, ValuesView +from typing import Any, TypeVar, overload + +from jax.tree_util import register_pytree_node_class + +_T = TypeVar("_T") +K = TypeVar("K") +V = TypeVar("V") + + +@register_pytree_node_class +class ImmutableMap(Mapping[K, V]): + """Immutable string-keyed dictionary. + + Parameters + ---------- + *args : tuple[str, V] + Key-value pairs. + **kwargs : V + Key-value pairs. + + Examples + -------- + >>> from galax.utils import ImmutableMap + >>> d = ImmutableMap(a=1, b=2) + >>> d + ImmutableMap({'a': 1, 'b': 2}) + """ + + def __init__( + self, + /, + *args: Mapping[K, V] | tuple[K, V] | Iterable[tuple[K, V]], + **kwargs: V, + ) -> None: + self._data: dict[K, V] = dict(*args, **kwargs) # type: ignore[assignment] + + # =========================================== + # Collection Protocol + + def __contains__(self, key: Any) -> bool: + return key in self._data + + def __iter__(self) -> Iterator[K]: + return iter(self._data) + + def __len__(self) -> int: + return len(self._data) + + # =========================================== + # Mapping Protocol + + def __getitem__(self, key: K) -> V: + return self._data[key] + + def keys(self) -> KeysView[K]: + return self._data.keys() + + def values(self) -> ValuesView[V]: + return self._data.values() + + def items(self) -> ItemsView[K, V]: + return self._data.items() + + @overload + def get(self, key: K, /) -> V | None: ... + + @overload + def get(self, key: K, /, default: V | _T) -> V | _T: ... + + def get(self, key: K, /, default: V | _T | None = None) -> V | _T | None: # type: ignore[misc] + return self._data.get(key, default) + + # =========================================== + # Extending Mapping + + def __or__(self, value: Any, /) -> "ImmutableMap[K, V]": + if not isinstance(value, Mapping): + return NotImplemented + + return type(self)(self._data | dict(value)) + + def __ror__(self, value: Any) -> Any: + return value | self._data + + # =========================================== + # Other + + def __hash__(self) -> int: + """Hash. + + Normally, dictionaries are not hashable because they are mutable. + However, this dictionary is immutable, so we can hash it. + """ + return hash(tuple(self._data.items())) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}({self._data!r})" + + # =========================================== + # JAX PyTree + + def tree_flatten(self) -> tuple[tuple[V, ...], tuple[K, ...]]: + """Flatten dict to the values (and keys). + + Returns + ------- + tuple[V, ...] tuple[str, ...] + A pair of an iterable with the values to be flattened recursively, + and the keys to pass back to the unflattening recipe. + """ + return (tuple(self._data.values()), tuple(self._data.keys())) + + @classmethod + def tree_unflatten( + cls, + aux_data: tuple[K, ...], + children: tuple[V, ...], + ) -> "ImmutableMap": # type: ignore[type-arg] # TODO: upstream beartype fix for ImmutableMap[V] + """Unflatten. + + Params: + aux_data: the opaque data that was specified during flattening of the + current treedef. + children: the unflattened children + + Returns + ------- + a re-constructed object of the registered type, using the specified + children and auxiliary data. + """ + return cls(tuple(zip(aux_data, children, strict=True))) diff --git a/src/immutable_map_jax/_version.pyi b/src/immutable_map_jax/_version.pyi new file mode 100644 index 0000000..91744f9 --- /dev/null +++ b/src/immutable_map_jax/_version.pyi @@ -0,0 +1,4 @@ +from __future__ import annotations + +version: str +version_tuple: tuple[int, int, int] | tuple[int, int, int, str, str] diff --git a/src/immutable_map_jax/py.typed b/src/immutable_map_jax/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_package.py b/tests/test_package.py new file mode 100644 index 0000000..7c72793 --- /dev/null +++ b/tests/test_package.py @@ -0,0 +1,9 @@ +from __future__ import annotations + +import importlib.metadata + +import immutable_map_jax as m + + +def test_version(): + assert importlib.metadata.version("immutable_map_jax") == m.__version__