Fix mypy

.github/workflows/ci.yml

@@ -39,6 +39,7 @@ jobs:
           report-title: Unit tests Linux - Python ${{ matrix.PYTHON_VERSION }}
 
   pre-commit-checks:
+    # TODO: switch to pixi once there is a good way
     name: pre-commit
     timeout-minutes: 30
     runs-on: ubuntu-latest
@@ -59,7 +60,7 @@ jobs:
         uses: pre-commit/[email protected]
 
   lint-workflow-files:
-    name: "Lint workflow files"
+    name: Lint workflow files
     runs-on: ubuntu-latest
     steps:
       - name: Checkout branch

pixi.toml

@@ -9,6 +9,7 @@ platforms = ["osx-arm64", "linux-64", "osx-64"]
 [tasks]
 "postinstall" = "pip install --no-build-isolation --no-deps --disable-pip-version-check -e ."
 "test" = "python polarify/test.py"
+"lint" = "pre-commit run --all"
 
 [dependencies]
 python = ">= 3.9"
@@ -20,3 +21,5 @@ python = ">= 3.9"
 "pytest" = "*"
 "pytest-md" = "*"
 "pytest-emoji" = "*"
+# linting
+"pre-commit" = "*"

polarify/__init__.py

@@ -9,12 +9,13 @@ def polarify(func):
     source = inspect.getsource(func)
     print(source)
     tree = ast.parse(source)
-    expr = parse_body(tree.body[0].body)
+    func_def: ast.FunctionDef = tree.body[0]  # type: ignore
+    expr = parse_body(func_def.body)
 
     # Replace the body of the function with the parsed expr
-    tree.body[0].body = [ast.Return(expr)]
+    func_def.body = [ast.Return(expr)]
     # TODO: make this prettier
-    tree.body[0].decorator_list = []
+    func_def.decorator_list = []
 
     # Unparse the modified AST back into source code
     new_func_code = ast.unparse(tree)