pyexec.jl

using PyCall

function pyeval_impl(source, globals, locals, lnn, mode)
    parse = pyimport("ast")[:parse]
    increment_lineno = pyimport("ast")[:increment_lineno]
    eval = pyimport("builtins")[:eval]
    compile = pyimport("builtins")[:compile]

    if lnn === nothing
        filename = "<PyCall>"
       lineno = 1
    else
        filename = lnn.file
        lineno = lnn.line + ('\n' in source ? 0 : -1)
    end

    source = rstrip(source)
    node = parse(source, filename, mode)
    increment_lineno(node, lineno)
    code = compile(node, filename, mode)
    eval(code, globals, locals)
end
# `rstrip(source)` is necessary to make the code like
#     @pyeval("""
#     [1,
#      2]
#      """)
# work.

"""
    @pyeval(source::String, [globals, [locals]])

Evaluate a Python expression in `source` in the context of `globals`
and `locals`.

This mimics Python's builtin `eval` function.  The main difference is
that `globals` always defaults to `__main__.__dict__`.  Note that, as
a consequence, `locals` defaults to it as well.

When writing multi-line expression, do not put code in the line with
`\"\"\"` (see also `@pyexec`):

```julia
@pyeval(\"\"\"
[1,
 2,
 3 // 0]
\"\"\")
```
"""
macro pyeval(source, globals=:($PyCall.maindict()), locals=nothing)
    lnn = source isa String ? QuoteNode(__source__) : nothing
    :(pyeval_impl($(esc(source)), $(esc(globals)), $(esc(locals)),
                  $lnn, "eval"))
end

"""
    @pyexec(source::String, [globals, [locals]])

Evaluate Python statements in `source` in the context of `globals` and
`locals`.

To get correct line numbers in the Python traceback, `source`
expression must start with `\"\"\"\\n`, e.g.:

```julia
@pyexec(\"\"\"
helper_const = object()

def helper_function():
    ...
\"\"\")
```

and not

```julia
@pyexec(\"\"\"helper_const = object()

           def helper_function():
               ...
        \"\"\")
```

We recommend not writing Python code in the line `\"\"\"` in general.
For example, the following code results in `IndentationError`.

```julia
@pyexec(\"\"\"def helper_function():
               ...\"\"\")
```

This mimics Python 3's builtin `exec` function.
"""
macro pyexec(source, globals=:($PyCall.maindict()), locals=nothing)
    lnn = source isa String ? QuoteNode(__source__) : nothing
    :(pyeval_impl($(esc(source)), $(esc(globals)), $(esc(locals)),
                  $lnn, "exec"))
end

function pyload_impl(source, globals, locals, __source__, mode)
    nstype = PyDict{Symbol, PyObject}
    if locals == nothing
        locals = nstype()
    elseif !(locals isa nstype)
        locals = PyDict(Dict{Symbol, PyObject}(locals))
    end
    if globals == nothing
        globals = locals
    elseif !(globals isa nstype)
        globals = PyDict(Dict{Symbol, PyObject}(globals))
    end
    pyeval_impl(source, globals, locals, __source__, mode)
    return locals
end

"""
    @pyload(source::String, [globals, [locals]]) :: PyDict

Evaluate Python statements in a fresh namespace and return the
`locals` namespace after the evaluation.  If only `globals` is
specified, it is used as the global namespace while an empty `locals`
is used for executing `source`.  If `locals` is not a `PyDict`, it is
converted to a `PyDict` first and returned after executing `source`.

To receive modification in `globals` namespace via `global` Python
statement, `globals` must be an instance of `PyDict{Symbol, PyObject}`.
"""
macro pyload(source, globals=nothing, locals=nothing)
    lnn = source isa String ? QuoteNode(__source__) : nothing
    :(pyload_impl($(esc(source)), $(esc(globals)), $(esc(locals)),
                  $lnn, "exec"))
end

example_eval() = @pyeval "1/0"

function example_exec()
    @pyexec """
    def f():
        1 / 0
    
    f()
    """
end