diff --git a/.github/actions/build-manylinux/entrypoint.sh b/.github/actions/build-manylinux/entrypoint.sh index 24fc59c4..20e57ba2 100755 --- a/.github/actions/build-manylinux/entrypoint.sh +++ b/.github/actions/build-manylinux/entrypoint.sh @@ -31,7 +31,7 @@ cmake --build build --config Release -- VERBOSE=1 export PATH=$PATH:$(pwd)/llvm/bin export LD_LIBRARY_PATH=$(pwd)/build:$LD_LIBRARY_PATH export CODON_DIR=$(pwd)/build -python3 -m pip install cython wheel +python3 -m pip install cython wheel astunparse python3 -m pip debug --verbose (cd extra/python; python3 setup.py sdist bdist_wheel --plat-name=manylinux2014_x86_64) python3 -m pip install -v extra/python/dist/*.whl diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c7ba4e41..5f2dcccd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -126,7 +126,7 @@ jobs: - name: Set up Python run: | python -m pip install --upgrade pip setuptools wheel - python -m pip install numpy cython wheel + python -m pip install numpy cython wheel astunparse which python which pip echo "CODON_PYTHON=$(python test/python/find-python-library.py)" >> $GITHUB_ENV diff --git a/codon/compiler/jit.cpp b/codon/compiler/jit.cpp index 28cbc763..c1c1d5a6 100644 --- a/codon/compiler/jit.cpp +++ b/codon/compiler/jit.cpp @@ -82,9 +82,11 @@ llvm::Error JIT::compile(const ir::Func *input) { return llvm::Error::success(); } -llvm::Expected JIT::compile(const std::string &code) { +llvm::Expected JIT::compile(const std::string &code, + const std::string &file, int line) { auto *cache = compiler->getCache(); - ast::StmtPtr node = ast::parseCode(cache, JIT_FILENAME, code, /*startLine=*/0); + ast::StmtPtr node = ast::parseCode(cache, file.empty() ? JIT_FILENAME : file, code, + /*startLine=*/line); auto sctx = cache->imports[MAIN_IMPORT].ctx; auto preamble = std::make_shared>(); @@ -166,8 +168,11 @@ llvm::Expected JIT::run(const ir::Func *input) { return getCapturedOutput(); } -llvm::Expected JIT::execute(const std::string &code) { - auto result = compile(code); +llvm::Expected +JIT::execute(const std::string &code, const std::string &file, int line, bool debug) { + if (debug) + fmt::print(stderr, "[codon::jit::execute] code:\n{}-----\n", code); + auto result = compile(code, file, line); if (auto err = result.takeError()) return std::move(err); if (auto err = compile(result.get())) @@ -198,7 +203,9 @@ std::string buildKey(const std::string &name, const std::vector &ty } std::string buildPythonWrapper(const std::string &name, const std::string &wrapname, - const std::vector &types) { + const std::vector &types, + const std::string &pyModule, + const std::vector &pyVars) { std::stringstream wrap; wrap << "@export\n"; wrap << "def " << wrapname << "(args: cobj) -> cobj:\n"; @@ -207,12 +214,21 @@ std::string buildPythonWrapper(const std::string &name, const std::string &wrapn << "a" << i << " = " << types[i] << ".__from_py__(PyTuple_GetItem(args, " << i << "))\n"; } + for (unsigned i = 0; i < pyVars.size(); i++) { + wrap << " " + << "py" << i << " = pyobj._get_module(\"" << pyModule << "\")._getattr(\"" + << pyVars[i] << "\")\n"; + } wrap << " return " << name << "("; for (unsigned i = 0; i < types.size(); i++) { if (i > 0) wrap << ", "; wrap << "a" << i; } + for (unsigned i = 0; i < pyVars.size(); i++) { + wrap << ", " + << "py" << i; + } wrap << ").__to_py__()\n"; return wrap.str(); @@ -228,8 +244,9 @@ ir::types::Type *JIT::PythonData::getCObjType(ir::Module *M) { return cobj; } -JITResult JIT::executeSafe(const std::string &code) { - auto result = execute(code); +JITResult JIT::executeSafe(const std::string &code, const std::string &file, int line, + bool debug) { + auto result = execute(code, file, line, debug); if (auto err = result.takeError()) { auto errorInfo = llvm::toString(std::move(err)); return JITResult::error(errorInfo); @@ -238,8 +255,10 @@ JITResult JIT::executeSafe(const std::string &code) { } JITResult JIT::executePython(const std::string &name, - const std::vector &types, void *arg) { - + const std::vector &types, + const std::string &pyModule, + const std::vector &pyVars, void *arg, + bool debug) { auto key = buildKey(name, types); auto &cache = pydata->cache; auto it = cache.find(key); @@ -253,7 +272,9 @@ JITResult JIT::executePython(const std::string &name, } else { static int idx = 0; auto wrapname = "__codon_wrapped__" + name + "_" + std::to_string(idx++); - auto wrapper = buildPythonWrapper(name, wrapname, types); + auto wrapper = buildPythonWrapper(name, wrapname, types, pyModule, pyVars); + if (debug) + fmt::print(stderr, "[codon::jit::executePython] wrapper:\n{}-----\n", wrapper); if (auto err = compile(wrapper).takeError()) { auto errorInfo = llvm::toString(std::move(err)); return JITResult::error(errorInfo); diff --git a/codon/compiler/jit.h b/codon/compiler/jit.h index 8da3e332..caed9146 100644 --- a/codon/compiler/jit.h +++ b/codon/compiler/jit.h @@ -51,18 +51,25 @@ public: // General llvm::Error init(); llvm::Error compile(const ir::Func *input); - llvm::Expected compile(const std::string &code); + llvm::Expected compile(const std::string &code, + const std::string &file = "", int line = 0); llvm::Expected address(const ir::Func *input); llvm::Expected run(const ir::Func *input); - llvm::Expected execute(const std::string &code); + llvm::Expected execute(const std::string &code, + const std::string &file = "", int line = 0, + bool debug = false); // Python llvm::Expected runPythonWrapper(const ir::Func *wrapper, void *arg); llvm::Expected getWrapperFunc(const std::string &name, const std::vector &types); JITResult executePython(const std::string &name, - const std::vector &types, void *arg); - JITResult executeSafe(const std::string &code); + const std::vector &types, + const std::string &pyModule, + const std::vector &pyVars, void *arg, + bool debug); + JITResult executeSafe(const std::string &code, const std::string &file, int line, + bool debug); // Errors llvm::Error handleJITError(const JITError &e); diff --git a/codon/parser/visitors/typecheck/op.cpp b/codon/parser/visitors/typecheck/op.cpp index 4aa90717..ca558a60 100644 --- a/codon/parser/visitors/typecheck/op.cpp +++ b/codon/parser/visitors/typecheck/op.cpp @@ -613,6 +613,13 @@ ExprPtr TypecheckVisitor::transformBinaryMagic(BinaryExpr *expr) { auto rt = expr->rexpr->getType()->getClass(); seqassert(lt && rt, "lhs and rhs types not known"); + if (!lt->is("pyobj") && rt->is("pyobj")) { + // Special case: `obj op pyobj` -> `rhs.__rmagic__(lhs)` on lhs + // Assumes that pyobj implements all left and right magics + return transform(N(N(expr->rexpr, format("__{}__", rightMagic)), + expr->lexpr)); + } + // Normal operations: check if `lhs.__magic__(lhs, rhs)` exists auto method = findBestMethod(lt, format("__{}__", magic), {lt, rt}); @@ -626,16 +633,6 @@ ExprPtr TypecheckVisitor::transformBinaryMagic(BinaryExpr *expr) { // Normal case: `__magic__(lhs, rhs)` return transform( N(N(method->ast->name), expr->lexpr, expr->rexpr)); - } else if (lt->is("pyobj")) { - // Special case: call `pyobj._getattr(magic)` on lhs - return transform(N(N(N(expr->lexpr, "_getattr"), - N(format("__{}__", magic))), - expr->rexpr)); - } else if (rt->is("pyobj")) { - // Special case: call `pyobj._getattr(magic)` on rhs - return transform(N(N(N(expr->rexpr, "_getattr"), - N(format("__r{}__", magic))), - expr->lexpr)); } return nullptr; } diff --git a/codon/parser/visitors/typecheck/typecheck.cpp b/codon/parser/visitors/typecheck/typecheck.cpp index 275ca9c6..eb0b68bf 100644 --- a/codon/parser/visitors/typecheck/typecheck.cpp +++ b/codon/parser/visitors/typecheck/typecheck.cpp @@ -250,7 +250,8 @@ bool TypecheckVisitor::wrapExpr(ExprPtr &expr, const TypePtr &expectedType, if (callee && expr->isType()) expr = transform(N(expr, N())); - std::unordered_set hints = {"Generator", "float", TYPE_OPTIONAL}; + std::unordered_set hints = {"Generator", "float", TYPE_OPTIONAL, + "pyobj"}; if (!exprClass && expectedClass && in(hints, expectedClass->name)) { return false; // argument type not yet known. } else if (expectedClass && expectedClass->name == "Generator" && @@ -267,6 +268,16 @@ bool TypecheckVisitor::wrapExpr(ExprPtr &expr, const TypePtr &expectedType, exprClass->name == TYPE_OPTIONAL && exprClass->name != expectedClass->name) { // unwrap optional expr = transform(N(N(FN_UNWRAP), expr)); + } else if (expectedClass && expectedClass->name == "pyobj" && + exprClass->name != expectedClass->name) { // wrap to pyobj + expr = transform( + N(N("pyobj"), N(N(expr, "__to_py__")))); + } else if (allowUnwrap && expectedClass && exprClass && exprClass->name == "pyobj" && + exprClass->name != expectedClass->name) { // unwrap pyobj + auto texpr = N(expectedClass->name); + texpr->setType(expectedType); + expr = + transform(N(N(texpr, "__from_py__"), N(expr, "p"))); } else if (callee && exprClass && expr->type->getFunc() && !(expectedClass && expectedClass->name == "Function")) { // Case 7: wrap raw Seq functions into Partial(...) call for easy realization. diff --git a/docs/interop/decorator.md b/docs/interop/decorator.md index 84e406e1..f654f7de 100644 --- a/docs/interop/decorator.md +++ b/docs/interop/decorator.md @@ -108,13 +108,68 @@ class Foo: def total(self): return self.a + self.b + self.c -print(Foo(10).total()) +print(Foo(10).total()) # 1110 ``` `@codon.convert` requires the annotated class to specify `__slots__`, which it uses to construct a generic Codon class (specifically, a named tuple) to store the class's converted fields. +# Passing globals to Codon + +Global variables, functions or modules can be passed to JIT'd functions through +the `pyvars` argument to `@codon.jit`: + +``` python +import codon + +def foo(n): + print(f'n is {n}') + +@codon.jit(pyvars=['foo']) +def bar(n): + foo(n) # calls the Python function 'foo' + return n ** 2 + +print(bar(9)) # 'n is 9' then '81' +``` + +This also allows imported Python modules to be accessed by Codon. All `pyvars` +are passed as Python objects. Note that JIT'd functions can call each other +by default. + +# Debugging + +`@codon.jit` takes an optional `debug` parameter that can be used to print debug +information such as generated Codon functions and data types: + +``` python +import codon + +@codon.jit(debug=True) +def sum_of_squares(v): + return sum(i**2 for i in v) + +print(sum_of_squares([1.4, 2.9, 3.14])) +``` + +outputs: + +``` +[codon::jit::execute] code: +def sum_of_squares(v): + return sum(i**2 for i in v) +----- +[python] sum_of_squares(['List[float]']) +[codon::jit::executePython] wrapper: +@export +def __codon_wrapped__sum_of_squares_0(args: cobj) -> cobj: + a0 = List[float].__from_py__(PyTuple_GetItem(args, 0)) + return sum_of_squares(a0).__to_py__() +----- +20.229599999999998 +``` + # Internals and performance tips Under the hood, the `codon` module maintains an instance of the Codon JIT, @@ -127,4 +182,6 @@ but instead reuses the cached function pointer. Although object conversions from Python to Codon are generally cheap, they do impose a small overhead, meaning **`@codon.jit` will work best on expensive and/or -long-running operations** rather than short-lived operations. +long-running operations** rather than short-lived operations. By the same token, +**the more work can be done in Codon, the better,** as opposed to repeatedly +transferring back and forth. diff --git a/docs/intro/releases.md b/docs/intro/releases.md index d1712137..f3272b33 100644 --- a/docs/intro/releases.md +++ b/docs/intro/releases.md @@ -96,7 +96,8 @@ for a in some_arbitrary_generator(): Included revamped `codon` module for Python, with `@codon.jit` decorator for compiling Python code in existing codebases. Further improved and -optimized the Python bridge. +optimized the Python bridge. Please see the [docs](../interop/decorator.md) +for more information. ## Codon IR diff --git a/extra/python/.gitignore b/extra/python/.gitignore index 7e7430b6..450b1aa7 100644 --- a/extra/python/.gitignore +++ b/extra/python/.gitignore @@ -1,2 +1,3 @@ dist -src/stdlib +codon/stdlib +codon/jit.cpp diff --git a/extra/python/MANIFEST.in b/extra/python/MANIFEST.in index 4c9e00e4..2543b799 100644 --- a/extra/python/MANIFEST.in +++ b/extra/python/MANIFEST.in @@ -1,3 +1,3 @@ -recursive-include src/stdlib *.codon -recursive-include src *.so -recursive-include src *.dylib +recursive-include codon/stdlib *.codon +recursive-include codon *.so +recursive-include codon *.dylib diff --git a/extra/python/README.md b/extra/python/README.md index 8c2878cb..e8ca5e76 100644 --- a/extra/python/README.md +++ b/extra/python/README.md @@ -7,8 +7,8 @@ $ pip install extra/python To use: ```python -from codon import codon, JitError +import codon -@codon +@codon.jit def ... ``` diff --git a/extra/python/src/__init__.py b/extra/python/codon/__init__.py similarity index 100% rename from extra/python/src/__init__.py rename to extra/python/codon/__init__.py diff --git a/extra/python/codon/decorator.py b/extra/python/codon/decorator.py new file mode 100644 index 00000000..5763729a --- /dev/null +++ b/extra/python/codon/decorator.py @@ -0,0 +1,176 @@ +from argparse import ArgumentError +import ctypes +import inspect +import sys +import os +import functools +import itertools +import ast +import astunparse + +sys.setdlopenflags(sys.getdlopenflags() | ctypes.RTLD_GLOBAL) + +if "CODON_PATH" not in os.environ: + os.environ["CODON_PATH"] = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) + os.environ["CODON_PATH"] += "/stdlib" + +from .codon_jit import JITWrapper, JITError + +pod_conversions = {type(None): "pyobj", + int: "int", + float: "float", + bool: "bool", + str: "str", + complex: "complex", + slice: "slice"} + +custom_conversions = {} +_error_msgs = set() + +def _common_type(t, debug, sample_size): + sub, is_optional = None, False + for i in itertools.islice(t, sample_size): + if i is None: + is_optional = True + else: + s = _codon_type(i, debug=debug, sample_size=sample_size) + if sub and sub != s: + return "pyobj" + sub = s + if is_optional and sub and sub != "pyobj": + sub = f"Optional[{sub}]" + return sub if sub else "pyobj" + +def _codon_type(arg, **kwargs): + t = type(arg) + + s = pod_conversions.get(t, "") + if s: + return s + if issubclass(t, list): + return f"List[{_common_type(arg, **kwargs)}]" + if issubclass(t, set): + return f"Set[{_common_type(arg, **kwargs)}]" + if issubclass(t, dict): + return f"Dict[{_common_type(arg.keys(), **kwargs)},{_common_type(arg.values(), **kwargs)}]" + if issubclass(t, tuple): + return f"Tuple[{','.join(_codon_type(a, **kwargs) for a in arg)}]" + s = custom_conversions.get(t, "") + if s: + j = ','.join(_codon_type(getattr(arg, slot), **kwargs) for slot in t.__slots__) + return f"{s}[{j}]" + + debug = kwargs.get('debug', None) + if debug: + msg = f"cannot convert {t.__name__}" + if msg not in _error_msgs: + print(f"[python] {msg}", file=sys.stderr) + _error_msgs.add(msg) + return "pyobj" + +def _codon_types(args, **kwargs): + return tuple(_codon_type(arg, **kwargs) for arg in args) + +def _reset_jit(): + global _jit + _jit = JITWrapper() + init_code = ("from internal.python import " + "setup_decorator, PyTuple_GetItem, PyObject_GetAttrString\n" + "setup_decorator()\n") + _jit.execute(init_code, "", 0, False) + return _jit + +_jit = _reset_jit() + +class RewriteFunctionArgs(ast.NodeTransformer): + def __init__(self, args): + self.args = args + + def visit_FunctionDef(self, node): + for a in self.args: + node.args.args.append(ast.arg(arg=a, annotation=None)) + return node + +def _obj_to_str(obj, **kwargs) -> str: + if inspect.isclass(obj): + lines = inspect.getsourcelines(obj)[0] + extra_spaces = lines[0].find("class") + obj_str = "".join(l[extra_spaces:] for l in lines) + elif callable(obj): + lines = inspect.getsourcelines(obj)[0] + extra_spaces = lines[0].find("@") + obj_str = "".join(l[extra_spaces:] for l in lines[1:]) + if kwargs.get('pyvars', None): + node = ast.fix_missing_locations( + RewriteFunctionArgs(kwargs['pyvars']).visit(ast.parse(obj_str)) + ) + obj_str = astunparse.unparse(node) + else: + raise TypeError(f"Function or class expected, got {type(obj).__name__}.") + return obj_str.replace("_@par", "@par") + +def _obj_name(obj) -> str: + if inspect.isclass(obj) or callable(obj): + return obj.__name__ + else: + raise TypeError(f"Function or class expected, got {type(obj).__name__}.") + +def _parse_decorated(obj, **kwargs): + return _obj_name(obj), _obj_to_str(obj, **kwargs) + +def convert(t): + if not hasattr(t, "__slots__"): + raise JITError(f"class '{str(t)}' does not have '__slots__' attribute") + + name = t.__name__ + slots = t.__slots__ + code = ("@tuple\n" + "class " + name + "[" + ",".join(f"T{i}" for i in range(len(slots))) + "]:\n") + for i, slot in enumerate(slots): + code += f" {slot}: T{i}\n" + + # PyObject_GetAttrString + code += " def __from_py__(p: cobj):\n" + for i, slot in enumerate(slots): + code += f" a{i} = T{i}.__from_py__(PyObject_GetAttrString(p, '{slot}'.ptr))\n" + code += f" return {name}({', '.join(f'a{i}' for i in range(len(slots)))})\n" + + _jit.execute(code, "", 0, False) + custom_conversions[t] = name + return t + +def jit(fn=None, debug=None, sample_size=5, pyvars=None): + if not pyvars: + pyvars = [] + if not isinstance(pyvars, list): + raise ArgumentError("pyvars must be a list") + def _decorate(f): + try: + obj_name, obj_str = _parse_decorated(f, pyvars=pyvars) + _jit.execute( + obj_str, + f.__code__.co_filename, + f.__code__.co_firstlineno, + 1 if debug else 0 + ) + except JITError: + _reset_jit() + raise + + @functools.wraps(f) + def wrapped(*args, **kwargs): + try: + args = (*args, *kwargs.values()) + types = _codon_types(args, debug=debug, sample_size=sample_size) + if debug: + print(f"[python] {f.__name__}({list(types)})", file=sys.stderr) + return _jit.run_wrapper(obj_name, types, f.__module__, pyvars, args, 1 if debug else 0) + except JITError: + _reset_jit() + raise + return wrapped + if fn: + return _decorate(fn) + return _decorate diff --git a/extra/python/src/jit.pxd b/extra/python/codon/jit.pxd similarity index 71% rename from extra/python/src/jit.pxd rename to extra/python/codon/jit.pxd index fa5019a2..c338c501 100644 --- a/extra/python/src/jit.pxd +++ b/extra/python/codon/jit.pxd @@ -14,5 +14,5 @@ cdef extern from "codon/compiler/jit.h" namespace "codon::jit": cdef cppclass JIT: JIT(string) Error init() - JITResult executeSafe(string) - JITResult executePython(string, vector[string], object) + JITResult executeSafe(string, string, int, char) + JITResult executePython(string, vector[string], string, vector[string], object, char) diff --git a/extra/python/src/jit.pyx b/extra/python/codon/jit.pyx similarity index 60% rename from extra/python/src/jit.pyx rename to extra/python/codon/jit.pyx index 4170c0f5..dce27511 100644 --- a/extra/python/src/jit.pyx +++ b/extra/python/codon/jit.pyx @@ -7,7 +7,7 @@ from cython.operator import dereference as dref from libcpp.string cimport string from libcpp.vector cimport vector -from src.jit cimport JIT, JITResult +from codon.jit cimport JIT, JITResult class JITError(Exception): @@ -24,16 +24,17 @@ cdef class JITWrapper: def __dealloc__(self): del self.jit - def execute(self, code: str) -> str: - result = dref(self.jit).executeSafe(code) + def execute(self, code: str, filename: str, fileno: int, debug: char) -> str: + result = dref(self.jit).executeSafe(code, filename, fileno, debug) if result: return None else: raise JITError(result.message) - def run_wrapper(self, name: str, types: list[str], args) -> object: + def run_wrapper(self, name: str, types: list[str], module: str, pyvars: list[str], args, debug: char) -> object: cdef vector[string] types_vec = types - result = dref(self.jit).executePython(name, types_vec, args) + cdef vector[string] pyvars_vec = pyvars + result = dref(self.jit).executePython(name, types_vec, module, pyvars_vec, args, debug) if result: return result.result else: diff --git a/extra/python/setup.py b/extra/python/setup.py index 730bd055..1a5057a9 100644 --- a/extra/python/setup.py +++ b/extra/python/setup.py @@ -49,11 +49,27 @@ codon_dir = Path(os.environ.get("CODON_DIR", from_root("build"))) codon_include_dir = os.environ.get("CODON_INCLUDE_DIR", codon_dir / "include") ext = "dylib" if sys.platform == "darwin" else "so" +def symlink(target, dest): + tmp = "_tmp" + os.symlink(str(target.resolve()), tmp) + os.rename(tmp, str(dest)) root = Path(os.path.dirname(os.path.realpath(__file__))) -distutils.dir_util.copy_tree(str(codon_dir / ".." / "stdlib"), str(root / "src" / "stdlib")) -shutil.copy(codon_dir / "lib" / "codon" / ("libcodonc." + ext), root / "src") -shutil.copy(codon_dir / "lib" / "codon" / ("libcodonrt." + ext), root / "src") -shutil.copy(codon_dir / "lib" / "codon" / ("libomp." + ext), root / "src") +symlink( + codon_dir / ".." / "stdlib", + root / "codon" / "stdlib" +) +symlink( + codon_dir / "lib" / "codon" / ("libcodonc." + ext), + root / "codon" / ("libcodonc." + ext) +) +symlink( + codon_dir / "lib" / "codon" / ("libcodonrt." + ext), + root / "codon" / ("libcodonrt." + ext) +) +symlink( + codon_dir / "lib" / "codon" / ("libomp." + ext), + root / "codon" / ("libomp." + ext) +) print(f" {llvm_include_dir}, {llvm_lib_dir}") print(f" {codon_include_dir}") @@ -65,18 +81,19 @@ else: jit_extension = Extension( "codon.codon_jit", - sources=["src/jit.pyx"], + sources=["codon/jit.pyx"], libraries=["codonc", "codonrt"], language="c++", extra_compile_args=["-w", "-std=c++17"], extra_link_args=[linker_args], include_dirs=[llvm_include_dir, str(codon_include_dir)], - library_dirs=[llvm_lib_dir, str(root / "src")], + library_dirs=[llvm_lib_dir, str(root / "codon")], ) setup( name="codon", version=CODON_VERSION, + install_requires=["astunparse"], python_requires='>=3.6', description="Codon JIT decorator", url="https://exaloop.io", @@ -87,6 +104,5 @@ setup( cmdclass={"build_ext": build_ext}, ext_modules=[jit_extension], packages=["codon"], - package_dir={"codon": "src"}, include_package_data=True ) diff --git a/extra/python/src/decorator.py b/extra/python/src/decorator.py deleted file mode 100644 index 860b11a8..00000000 --- a/extra/python/src/decorator.py +++ /dev/null @@ -1,141 +0,0 @@ -import ctypes -import inspect -import importlib -import importlib.util -import sys -import os - -sys.setdlopenflags(sys.getdlopenflags() | ctypes.RTLD_GLOBAL) - -if "CODON_PATH" not in os.environ: - os.environ["CODON_PATH"] = os.path.dirname( - os.path.abspath(inspect.getfile(inspect.currentframe())) - ) - os.environ["CODON_PATH"] += "/stdlib" - -from .codon_jit import JITWrapper, JITError - -pod_conversions = {type(None): "NoneType", - int: "int", - float: "float", - bool: "bool", - str: "str", - complex: "complex", - slice: "slice"} - -custom_conversions = {} - -def _codon_type(arg): - t = type(arg) - - s = pod_conversions.get(t, "") - if s: - return s - - if issubclass(t, list): - sub = "NoneType" - x = next(iter(arg), None) - if x is not None: - sub = _codon_type(x) - return f"List[{sub}]" - - if issubclass(t, set): - sub = "NoneType" - x = next(iter(arg), None) - if x is not None: - sub = _codon_type(x) - return f"Set[{sub}]" - - if issubclass(t, dict): - sub1 = "NoneType" - sub2 = "NoneType" - x = next(iter(arg.items()), None) - if x is not None: - sub1 = _codon_type(x[0]) - sub2 = _codon_type(x[1]) - return f"Dict[{sub1},{sub2}]" - - if issubclass(t, tuple): - return f"Tuple[{','.join(_codon_type(a) for a in arg)}]" - - s = custom_conversions.get(t, "") - if s: - return f"{s}[{','.join(_codon_type(getattr(arg, slot)) for slot in t.__slots__)}]" - - return "pyobj" - -def _codon_types(args): - return tuple(_codon_type(arg) for arg in args) - -def _reset_jit(): - global _jit - _jit = JITWrapper() - init_code = ("from internal.python import " - "setup_decorator, PyTuple_GetItem, PyObject_GetAttrString\n" - "setup_decorator()\n") - _jit.execute(init_code) - return _jit - -_jit = _reset_jit() - -def _obj_to_str(obj) -> str: - if inspect.isclass(obj): - lines = inspect.getsourcelines(obj)[0] - extra_spaces = lines[0].find("class") - obj_str = "".join(l[extra_spaces:] for l in lines) - elif callable(obj): - lines = inspect.getsourcelines(obj)[0] - extra_spaces = lines[0].find("@") - obj_str = "".join(l[extra_spaces:] for l in lines[1:]) - else: - raise TypeError(f"Function or class expected, got {type(obj).__name__}.") - return obj_str.replace("_@par", "@par") - -def _obj_name(obj) -> str: - if inspect.isclass(obj) or callable(obj): - return obj.__name__ - else: - raise TypeError(f"Function or class expected, got {type(obj).__name__}.") - -def _parse_decorated(obj): - return _obj_name(obj), _obj_to_str(obj) - -def convert(t): - if not hasattr(t, "__slots__"): - raise JITError(f"class '{str(t)}' does not have '__slots__' attribute") - - name = t.__name__ - slots = t.__slots__ - code = ("@tuple\n" - "class " + name + "[" + ",".join(f"T{i}" for i in range(len(slots))) + "]:\n") - for i, slot in enumerate(slots): - code += f" {slot}: T{i}\n" - - # PyObject_GetAttrString - code += " def __from_py__(p: cobj):\n" - for i, slot in enumerate(slots): - code += f" a{i} = T{i}.__from_py__(PyObject_GetAttrString(p, '{slot}'.ptr))\n" - code += f" return {name}({', '.join(f'a{i}' for i in range(len(slots)))})\n" - - _jit.execute(code) - custom_conversions[t] = name - return t - -def jit(obj): - try: - obj_name, obj_str = _parse_decorated(obj) - _jit.execute(obj_str) - except JITError as e: - _reset_jit() - raise - - def wrapped(*args, **kwargs): - try: - args = (*args, *kwargs.values()) - types = _codon_types(args) - return _jit.run_wrapper(obj_name, types, args) - except JITError as e: - _reset_jit() - raise - - return wrapped diff --git a/stdlib/internal/python.codon b/stdlib/internal/python.codon index f882d1a2..26acae30 100644 --- a/stdlib/internal/python.codon +++ b/stdlib/internal/python.codon @@ -9,9 +9,11 @@ Py_DecRef = Function[[cobj], NoneType](cobj()) Py_IncRef = Function[[cobj], NoneType](cobj()) Py_Initialize = Function[[], NoneType](cobj()) PyImport_AddModule = Function[[cobj], cobj](cobj()) +PyImport_AddModuleObject = Function[[cobj], cobj](cobj()) PyImport_ImportModule = Function[[cobj], cobj](cobj()) PyErr_Fetch = Function[[Ptr[cobj], Ptr[cobj], Ptr[cobj]], NoneType](cobj()) PyRun_SimpleString = Function[[cobj], NoneType](cobj()) +PyEval_GetGlobals = Function[[], cobj](cobj()) # conversions PyLong_AsLong = Function[[cobj], int](cobj()) @@ -25,6 +27,7 @@ PyList_GetItem = Function[[cobj, int], cobj](cobj()) PyList_SetItem = Function[[cobj, int, cobj], cobj](cobj()) PyDict_New = Function[[], cobj](cobj()) PyDict_Next = Function[[cobj, Ptr[int], Ptr[cobj], Ptr[cobj]], int](cobj()) +PyDict_GetItem = Function[[cobj, cobj], cobj](cobj()) PyDict_SetItem = Function[[cobj, cobj, cobj], cobj](cobj()) PySet_Add = Function[[cobj, cobj], cobj](cobj()) PySet_New = Function[[cobj], cobj](cobj()) @@ -148,9 +151,11 @@ def init_dl_handles(py_handle: cobj): global Py_IncRef global Py_Initialize global PyImport_AddModule + global PyImport_AddModuleObject global PyImport_ImportModule global PyErr_Fetch global PyRun_SimpleString + global PyEval_GetGlobals global PyLong_AsLong global PyLong_FromLong global PyFloat_AsDouble @@ -162,6 +167,7 @@ def init_dl_handles(py_handle: cobj): global PyList_SetItem global PyDict_New global PyDict_Next + global PyDict_GetItem global PyDict_SetItem global PySet_Add global PySet_New @@ -235,9 +241,11 @@ def init_dl_handles(py_handle: cobj): Py_IncRef = dlsym(py_handle, "Py_IncRef") Py_Initialize = dlsym(py_handle, "Py_Initialize") PyImport_AddModule = dlsym(py_handle, "PyImport_AddModule") + PyImport_AddModuleObject = dlsym(py_handle, "PyImport_AddModuleObject") PyImport_ImportModule = dlsym(py_handle, "PyImport_ImportModule") PyErr_Fetch = dlsym(py_handle, "PyErr_Fetch") PyRun_SimpleString = dlsym(py_handle, "PyRun_SimpleString") + PyEval_GetGlobals = dlsym(py_handle, "PyEval_GetGlobals") PyLong_AsLong = dlsym(py_handle, "PyLong_AsLong") PyLong_FromLong = dlsym(py_handle, "PyLong_FromLong") PyFloat_AsDouble = dlsym(py_handle, "PyFloat_AsDouble") @@ -249,6 +257,7 @@ def init_dl_handles(py_handle: cobj): PyList_SetItem = dlsym(py_handle, "PyList_SetItem") PyDict_New = dlsym(py_handle, "PyDict_New") PyDict_Next = dlsym(py_handle, "PyDict_Next") + PyDict_GetItem = dlsym(py_handle, "PyDict_GetItem") PyDict_SetItem = dlsym(py_handle, "PyDict_SetItem") PySet_Add = dlsym(py_handle, "PySet_Add") PySet_New = dlsym(py_handle, "PySet_New") @@ -579,6 +588,9 @@ class pyobj: pyobj.decref(it) pyobj.exc_check() + def to_str(self, errors: str, empty: str = "") -> str: + return pyobj.to_str(self.p, errors, empty) + def to_str(p: cobj, errors: str, empty: str = "") -> str: obj = PyUnicode_AsEncodedString(p, "utf-8".c_str(), errors.c_str()) if obj == cobj(): @@ -588,9 +600,6 @@ class pyobj: pyobj.decref(obj) return res - def to_str(self, errors: str, empty: str = "") -> str: - return pyobj.to_str(self.p, errors, empty) - def exc_check(): ptype, pvalue, ptraceback = cobj(), cobj(), cobj() PyErr_Fetch(__ptr__(ptype), __ptr__(pvalue), __ptr__(ptraceback)) @@ -657,8 +666,15 @@ class pyobj: ensure_initialized() PyRun_SimpleString(code.c_str()) + def _globals() -> pyobj: + return pyobj(PyEval_GetGlobals()) + + def _get_module(name: str) -> pyobj: + p = pyobj(pyobj.exc_wrap(PyImport_AddModule(name.c_str()))) + return p + def _main_module() -> pyobj: - return pyobj(pyobj.exc_wrap(PyImport_AddModule("__main__".c_str()))) + return pyobj._get_module("__main__") def _repr_mimebundle_(self, bundle=Set[str]()) -> Dict[str, str]: fn = pyobj._main_module()._getattr("__codon_repr__") diff --git a/stdlib/internal/types/float.codon b/stdlib/internal/types/float.codon index 72e0a0b3..c2e56d2e 100644 --- a/stdlib/internal/types/float.codon +++ b/stdlib/internal/types/float.codon @@ -16,10 +16,7 @@ class float: def __new__() -> float: return 0.0 - def __new__(what) -> float: - return what.__float__() - - def __new__(s: str) -> float: + def _from_str(s: str) -> float: from C import strtod(cobj, Ptr[cobj]) -> float buf = __array__[byte](32) @@ -41,6 +38,13 @@ class float: return result + def __new__(what) -> float: + # do not overload! (needed to avoid pyobj conversion) + if isinstance(what, str): + return float._from_str(what) + else: + return what.__float__() + def __repr__(self) -> str: s = seq_str_float(self) return s if s != "-nan" else "nan" diff --git a/stdlib/internal/types/int.codon b/stdlib/internal/types/int.codon index 8b24ffd2..b616d69c 100644 --- a/stdlib/internal/types/int.codon +++ b/stdlib/internal/types/int.codon @@ -23,15 +23,16 @@ class int: def __new__() -> int: ret i64 0 - def __new__(what) -> int: - return what.__int__() - - def __new__(s: str) -> int: - return int._from_str(s, 10) - def __new__(s: str, base: int) -> int: return int._from_str(s, base) + def __new__(what) -> int: + # do not overload! (needed to avoid pyobj conversion) + if isinstance(what, str): + return int._from_str(what, 10) + else: + return what.__int__() + def __int__(self) -> int: return self diff --git a/test/parser/simplify_stmt.codon b/test/parser/simplify_stmt.codon index 4e63632d..c3826e2b 100644 --- a/test/parser/simplify_stmt.codon +++ b/test/parser/simplify_stmt.codon @@ -917,6 +917,23 @@ def foo(): return str(pyfoo3()) print foo() #: 3 +#%% python_pyobj +@python +def foofn() -> Dict[pyobj, pyobj]: + return {"str": "hai", "int": 1} + +foo = foofn() +print(foo, foo.__class__.__name__) +#: {'str': 'hai', 'int': 1} Dict[pyobj,pyobj] +foo["codon"] = 5.15 +print(foo["codon"], foo["codon"].__class__.__name__, foo.__class__.__name__) +#: 5.15 pyobj Dict[pyobj,pyobj] + +a = {1: "s", 2: "t"} +a[3] = foo["str"] +print(a) +#: {1: 's', 2: 't', 3: 'hai'} + #%% typeof_definition_error,barebones a = 1 class X: diff --git a/test/python/cython_jit.py b/test/python/cython_jit.py index b8eb5f5d..499e109e 100644 --- a/test/python/cython_jit.py +++ b/test/python/cython_jit.py @@ -1,5 +1,3 @@ -import sys -from io import StringIO from typing import Dict, List, Tuple import codon @@ -45,7 +43,7 @@ def test_roundtrip(): def roundtrip(x): return x - for i in range(5): + for _ in range(5): assert roundtrip(None) == None assert roundtrip(42) == 42 assert roundtrip(3.14) == 3.14 @@ -113,3 +111,43 @@ test_roundtrip() test_return_type() test_param_types() test_error_handling() + + +@codon.jit +def foo(y): + return f"{y.__class__.__name__}; {y}" + + +@codon.jit(debug=True) +def foo2(y): + return f"{y.__class__.__name__}; {y}" + +class Foo: + def __init__(self): + self.x = 1 + +@codon.jit +def a(x): + return x+1 + +def b(x, z): + y = a(x) + return y * z + +@codon.jit(pyvars=['b']) +def c(x, y): + n = b(x,y) ** a(1) + return n + +def test_cross_calls(): + assert foo([None, 1]) == "List[Optional[int]]; [None, 1]" + assert foo([1, None, 1]) == "List[Optional[int]]; [1, None, 1]" + assert foo([1, None, 1.2]) == "List[pyobj]; [1, None, 1.2]" + assert foo({None: 1}) == "Dict[pyobj,int]; {None: 1}" + assert foo2([None, Foo()]).startswith("List[pyobj]; [None, <__main__.Foo object at") + + assert a(3) == 4 + assert b(3, 4) == 16 + assert round(c(5, 6.1), 2) == 1339.56 + +test_cross_calls()