migrate to uv (#54)

* uv initial support

* building

* linter fix

* pre-commit update

* fix typings

Author: Howuhh

.github/workflows/codestyle.yml

@@ -12,19 +12,30 @@ on:
 jobs:
   build:
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version:
+          - "3.9"
+          - "3.10"
+          - "3.11"
+          - "3.12"
     steps:
-      - uses: actions/checkout@v3
-      - name: Set up Python 3.10
-        uses: actions/setup-python@v4
+      - uses: actions/checkout@v5
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
         with:
-          python-version: "3.10"
+          version: "0.9.17"
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+
+      - name: Set up Python
+        run: uv python install
+
       - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install -e ".[dev]"
+        run: uv sync --locked --all-extras --dev
       - name: check codestyle
         run: |
-          ruff check --config pyproject.toml --diff .
+          uv run ruff check --config pyproject.toml --diff .
       - name: check type hints
         run: |
-          pyright --project=pyproject.toml src/xminigrid
+          uv run pyright --project=pyproject.toml

.gitignore

@@ -175,3 +175,6 @@ src/xminigrid/envs/xland_tmp.py
 scripts/*testing*
 configs/wip
 scripts/eval_example.py
+
+.ruff_cache/
+.pyright/

.pre-commit-config.yaml

@@ -1,17 +1,27 @@
+default_stages: [pre-commit]
+fail_fast: false
+
 repos:
-  # ruff checking
-  - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.6.0
+  - repo: local
     hooks:
-      # Run the linter.
       - id: ruff
-        args: [--fix]
-      # Run the formatter.
+        name: ruff (lint)
+        entry: uv run ruff check --fix
+        language: system
+        types_or: [python, pyi]
+        require_serial: true
+
       - id: ruff-format
+        name: ruff (format)
+        entry: uv run ruff format
+        language: system
+        types_or: [python, pyi]
+        require_serial: true
 
-  # pyright checking
-  - repo: https://github.com/RobertCraigie/pyright-python
-    rev: v1.1.376
-    hooks:
       - id: pyright
-        args: [--project=pyproject.toml]
+        name: pyright
+        entry: uv run pyright --project=pyproject.toml
+        language: system
+        types_or: [python, pyi]
+        require_serial: true
+        pass_filenames: false

CONTRIBUTING.md

@@ -13,7 +13,7 @@ Do not be afraid to ask questions and share new ideas about possible directions
 
 ## Contributing to the codebase
 
-Contributing code is done through standard github methods:
+We use uv for dependency management and virtual environments. Contributing code is done through standard github methods:
 
 1. Fork this repo
 2. Make a change and commit your code
@@ -22,8 +22,22 @@ Contributing code is done through standard github methods:
 ```commandline
 git clone [email protected]:corl-team/xland-minigrid.git
 cd xland-minigrid
-pip install -e ".[dev]"
+
+# Create and sync the development environment
+uv venv
+uv sync --extra dev --extra baselines
+```
+
+Add new dependencies using:
+
 ```
+uv add <package>
+
+# development-only dependencies
+uv add <package> --group dev
+```
+
+Do not use pip install for project dependencies, as it bypasses the lockfile and may cause inconsistent environments.
 
 ## Code style
 
@@ -32,17 +46,17 @@ The CI will run several checks on the new code pushed to the repository.
 These checks can also be run locally without waiting for the CI by following the steps below: 
 
 1. install [pre-commit](https://pre-commit.com/#install)
-2. install the Git hooks by running `pre-commit install`
+2. install the Git hooks by running `uv run pre-commit install`
 
 Once those two steps are done, the Git hooks will be run automatically at
 every new commit. The Git hooks can also be run manually with 
-`pre-commit run --all-files`, and if needed they can be 
+`uv run pre-commit run`, and if needed they can be 
 skipped (not recommended) with `git commit --no-verify`.
 
-Be sure to run and fix all issues from the `pre-commit run --all-files` before the push!
-If you want to see possible problems before pre-commit, you can run `ruff check --diff .` 
-and `ruff format --check` to see exact linter and formatter suggestions and possible fixes. 
-Similarly, run `pyright src/xminigrid` to see possible problems with type hints.
+Be sure to run and fix all issues from the `uv run pre-commit run --all-files` before the push!
+If you want to see possible problems before pre-commit, you can run `uv run ruff check --diff .` 
+and `uv run ruff format --check` to see exact linter and formatter suggestions and possible fixes. 
+Similarly, run `uv run pyright src/xminigrid` to see possible problems with type hints.
 
 # License
 

README.md

@@ -67,9 +67,9 @@ How cool is that? For more details, take a look at the [technical paper](https:/
 The latest release of XLand-MiniGrid can be installed directly from PyPI:
 
 ```commandline
-pip install xminigrid
+uv add xminigrid (or pip install xminigrid)
 # or, from github directly
-pip install "xminigrid @ git+https://github.com/corl-team/xland-minigrid.git"
+uv pip install "xminigrid @ git+https://github.com/corl-team/xland-minigrid.git"
 ```
 
 Alternatively, if you want to install the latest development version from the GitHub and run provided algorithms or scripts,
@@ -78,8 +78,8 @@ install the source as follows:
 git clone [email protected]:corl-team/xland-minigrid.git
 cd xland-minigrid
 
-# additional dependencies for baselines
-pip install -e ".[dev,baselines]"
+# additional dependencies for dev and baselines
+uv sync --all-extras --dev
 ```
 Note that the installation of JAX may differ depending on your hardware accelerator! 
 We advise users to explicitly install the correct JAX version (see the [official installation guide](https://github.com/google/jax#installation)).

pyproject.toml

@@ -5,7 +5,7 @@ readme = "README.md"
 requires-python =">=3.9"
 license = {file = "LICENSE"}
 authors = [
-  {name = "Alexander Nikulin", email = "[email protected]"},
+  {name = "Alexander Nikulin", email = "[email protected]"},
 ]
 dynamic = ["version"]
 
@@ -62,9 +62,16 @@ baselines = [
 
 
 [build-system]
-requires = ["setuptools>=61.0"]
+requires = ["setuptools>=61.0", "wheel"]
 build-backend = "setuptools.build_meta"
 
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
 [tool.setuptools.dynamic]
 version = {attr = "xminigrid.__version__"}
 
@@ -96,9 +103,13 @@ known-third-party = ["wandb"]
 include = ["src/xminigrid"]
 exclude = [
     'training',
+    'scripts',
     "**/node_modules",
     "**/__pycache__",
 ]
+executionEnvironments = [
+  { root = "src" }
+]
 reportMissingImports = "none"
 reportMissingTypeStubs = false
 reportMissingModuleSource = false

setup.py

@@ -2,7 +2,7 @@
 from .registration import make, register, registered_environments
 
 # TODO: add __all__
-__version__ = "0.9.1"
+__version__ = "0.9.2"
 
 # ---------- XLand-MiniGrid environments ----------
 

src/xminigrid/__init__.py

@@ -44,17 +44,16 @@ def check_goal(
 # should I inherit from abc.ABC?
 class BaseGoal(struct.PyTreeNode):
     @abc.abstractmethod
-    def __call__(self, grid: GridState, agent: AgentState, action: int | jax.Array, position: jax.Array) -> jax.Array:
-        ...
+    def __call__(
+        self, grid: GridState, agent: AgentState, action: int | jax.Array, position: jax.Array
+    ) -> jax.Array: ...
 
     @classmethod
     @abc.abstractmethod
-    def decode(cls, encoding: jax.Array) -> BaseGoal:
-        ...
+    def decode(cls, encoding: jax.Array) -> BaseGoal: ...
 
     @abc.abstractmethod
-    def encode(self) -> jax.Array:
-        ...
+    def encode(self) -> jax.Array: ...
 
 
 class EmptyGoal(BaseGoal):

src/xminigrid/core/goals.py

@@ -52,17 +52,14 @@ class BaseRule(struct.PyTreeNode):
     @abc.abstractmethod
     def __call__(
         self, grid: GridState, agent: AgentState, action: int | jax.Array, position: jax.Array
-    ) -> tuple[GridState, AgentState]:
-        ...
+    ) -> tuple[GridState, AgentState]: ...
 
     @classmethod
     @abc.abstractmethod
-    def decode(cls, encoding: jax.Array) -> BaseRule:
-        ...
+    def decode(cls, encoding: jax.Array) -> BaseRule: ...
 
     @abc.abstractmethod
-    def encode(self) -> jax.Array:
-        ...
+    def encode(self) -> jax.Array: ...
 
 
 # TODO: checks on only some actions can produce bugs, but it is faster. If it's going to cause trouble, reconsider.