Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

chore: update python tests to use new parser #9

Merged
merged 1 commit into from
Dec 23, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions NKL/FFI.lean
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,6 @@ def parse_json_old (json : String) : IO Unit := do
@[export parse_json]
def parse_json (s : String) : IO Unit := do
let kernel <- Python.Parsing.parse s
for (n,f) in kernel.funcs do
IO.println s!"found {n}"
IO.println s!"{repr f}"
let names := kernel.funcs.map fun x => x.fst
let names := String.intercalate "," names
IO.println s!"Found functions: {names}"
27 changes: 27 additions & 0 deletions NKL/Python.lean
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@ inductive Const where
| ellipsis
deriving Repr

-- We don't need these, but we preserve them to make copying
-- the Python AST over easier.
inductive Ctx where
| load | store | del
deriving Repr
Expand Down Expand Up @@ -80,6 +82,25 @@ inductive Stmt' where
deriving Repr
end

/-
This structure is a mirror of the python arguments AST node.
If we have the following Python function:

def f(a, b=1, /, c=2, *args, d, e=3, **kwargs): pass

then the structure will be populated with:

posonlyargs = [a, b]
args = [c]
defaults = [1, 2]
vararg = "args"
kwonlyargs = [d, e]
kw_defaults = [None, 3]
kwarg = "kwargs"

Note, defaults and kw_defaults are inconsistent in how they treat
missing arguments, but this is just how it works in the python AST.
-/
structure Args where
posonlyargs : List String
args : List String
Expand All @@ -90,6 +111,12 @@ structure Args where
kwarg : Option String
deriving Repr

/-
In addition to the defaults above from the AST, we also collect
the values from f.__defaults__ here in the Fun structure. These
values are evaluated in a different context from the other names
in the function, so we need to capture them on the Python side.
-/
structure Fun where
source : String
args : Args
Expand Down
23 changes: 23 additions & 0 deletions interop/nkl/parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,25 @@
from textwrap import dedent
from itertools import chain
from collections import deque
from nkl.lean import py_to_lean

# This is a custom JSON encoder for use with AST nodes.
# The AST nodes are not handled by the default encoder.
# For an AST node, we return a dictionary with the class
# name mapped to the object dictionary. If the object
# dictionary is empty we just return the class name.
# e.g.
# Binop(left=l,op=o,right=r), becomes:
# { BinOp : { left:l, op:o, right:r } }
# Pass(), becomes
# 'Pass'
#
# For anything else not handled by the default
# encoder, we return "...", the Ellipsis.
# Conveniently, Ellipsis is one of the things
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Won't we need ellipses for indexing?

# that isn't handled, so it is properly mapped.

# See also: NKL/Python.lean for the Lean side

class Enc(json.JSONEncoder):
def default(self, obj):
Expand Down Expand Up @@ -40,6 +59,10 @@ def json(self):
}
return json.dumps(d, cls=Enc)

# TODO: just a placeholder for testing
def load(self):
py_to_lean(self.json())

# resolve a reference: either populating the environment,
# or adding new items to the work queue
def reference(self, refname, val):
Expand Down
7 changes: 3 additions & 4 deletions interop/test/test_examples.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import pytest

from examples import *
from nkl.loader import Loader
from nkl.parser import Parser

@pytest.mark.parametrize("module,name",
[ (getting_started, "nki_tensor_add_kernel"),
Expand Down Expand Up @@ -31,6 +31,5 @@
])
def test_parse(module, name):
f = getattr(module, name)
F = Loader(f)
F.translate(F.ast)

F = Parser(f)
F.load()
50 changes: 0 additions & 50 deletions interop/test/test_json.py

This file was deleted.

Loading