Feature Request: streamline importing python code and type annotations

[strangemonad]

this proposal isn't completely thought-out w.r.t. type checking but there's at least some rough thoughts

Overview

This aims to streamline the core experience of using imported python code

1. The syntax for importing python modules
2. Using imported python types as type annotations

Dependencies

• This is likely going to have a dependency on the lifetime and borrow checking
• This is likely going to have an impact on type checking and implementation of co/contra variant generic types.

Proposal

There's 2 parts to the proposal

1. Allow importing python to produce module-level symbols
2. Allow using imported python symbols as type annotations where appropriate

1. Import syntax

Allow an alternative import syntax to directly import python types. I could be convinced that this special syntax isn't even needed. For the time being though, being explicit here seems prudent.

import python <module-spec> [as <ident>]

# or
from python <module-spec> import <ident> [as <ident>]

This would be roughly equivalent to let <ident> = Python.import_module("<module_name>").

• The declared value shares the same namespace scope as mojo imports
• Until the borrow checker is richer, it's probably fine to artificially extend the lifetime of imports indefinitely

2. Use imported python types as type annotations

Allow type annotations to match what you'd expect when annotating python

e.g.

import python numpy as np

def compute_something(x: np.array):
    pass

Motivating example

Let's use a simple example of benchmarking numpy's. In some cases you have to either thread imported python values as arguments or re-import them where you're using them.

from Benchmark import Benchmark
from PythonInterface import Python
from PythonObject import PythonObject

let np = Python.import_module("numpy")


fn periodic_signal(np: PythonObject, sampling_rate: Int) raises -> PythonObject:
    """Generates a single period periodic signal of length `sampling_rate`.
    
    Returns
    -------
        A numpy.ndarray with `sampling_rate` number of values
    
    """
    
    let ts = 1.0 / sampling_rate
    let t = np.arange(0, 1, ts)
    
    let frequency = 1.0
    let x = 3 * np.sin(2 * np.pi * frequency * t)
    return x


fn numpy_fft(np: PythonObject, input: PythonObject) -> None:
    try:
        _ = np.fft.fft(input)
    except:
        print("ignoring numpy fft error")


fn benchmark_numpy(input: PythonObject) raises:
    let np = Python.import_module("numpy")

    @parameter
    fn bench():
        _ = numpy_fft(np, input)

    let nanos = Benchmark(1000, 100000).run[bench]()
    
    print("numpy.fft.fft ns/op:")
    print(nanos)
    
    # XXX Artificially force a longer lifetime for closure-captured variables
    # until the lifetime checker improves.
    _ = np^

input = periodic_signal(np, 1024)
benchmark_numpy(input)

Things I haven't thought through deeply

• Dynamic module imports can't be compile time parameters, how does this gracefully degrade
• Are module imports that are nested to some scope shorter lived? e.g. if I import a module in a function, does it fall out of scope?
  • I'd argue that it shouldn't. Python's import mechanism assumes imported modules live for the lifetime of that cpython session.
• I haven't thought through the full type checking implications, I could imagine either
  • a type annotation like np.array becoming PythonObject[<NumpyArrayPyType>] where NumpyArrayPyType is the specific type. I could imagine this could even have a runtime representation by calling into the CPython API to get the PyTypeObject? This approach could also support generics as multiple parameters.
  • alternatively, you could represent NumpyArrayPyType as having a PythonObject upper bound.
  • I've not through through co/contra variance very deeply here. This would depend on if there are any flavor of in/out/inout parameters
• This obviously forces some amount of compile time evaluation of python module imports (but checkers like mypy or pyright already have to do this). It's unclear if this phase should match 1-1 with whatever const-expr phase is allowed for alias or @parameter expressions or if this warrants different semantics.