[Enhancement] [Python-Mojo bindings] Loosen `PythonObject` requirement for return type and arguments

Bug description

Actual behavior

The file provided below compiles correctly but the the resulting library doesn't expose any method from the mojo object.

Expected behavior

The object should expose all methods from the mojo compiled library.

Steps to reproduce

mojo file

from python import PythonObject
from python.bindings import PythonModuleBuilder
from os import abort
import math
from python import Python


@export
fn PyInit_odes() -> PythonObject:

    try:
        var m = PythonModuleBuilder("odes")
                # Expose the MFCModel struct as a Python class.
        _ = m.add_type[MFCModel]("MFCModel")

        # Expose the mfc_odes method on the MFCModel class.


        
        return m.finalize()

    except :
        print("Error creating MFCModel:")
        return PythonObject(None)


struct MFCModel( Defaultable, Movable, Representable, Copyable):
    """
    Microbial Fuel Cell (MFC) model parameters and state variables.
    (July 2025 API - Corrected Nested Syntax).
    """
    var F: Float64
    var R: Float64
    var T: Float64
    var k_m: Float64
    var d_m: Float64
    var k_aq: Float64
    var d_cell: Float64
    var C_a: Float64
    var C_c: Float64
    var V_a: Float64
    var V_c: Float64
    var A_m: Float64
    var Y_ac: Float64
    var K_dec: Float64
    var f_x: Float64
    var alpha: Float64
    var beta: Float64
    var k1_0: Float64
    var k2_0: Float64
    var K_AC: Float64
    var K_O2: Float64
    var Q_a: Float64        
    var Q_c: Float64
    var C_AC_in: Float64
    var C_CO2_in: Float64
    var X_in: Float64
    var C_H_in: Float64
    var C_O2_in: Float64
    var C_M_in: Float64
    var C_OH_in: Float64
    var U0: Float64
    var C_AC: Float64
    var C_CO2: Float64
    var C_H: Float64
    var X: Float64
    var C_O2: Float64
    var C_OH: Float64
    var C_M: Float64
    var eta_a: Float64
    var eta_c: Float64

        


    fn __init__(out self):
        """Initializes with default parameters using correct Mojo syntax."""
        self.F = 96485.332
        self.R = 8.314
        self.T = 303
        self.k_m = 17.0
        self.d_m = 1.778e-4
        self.k_aq = 5.0
        self.d_cell = 2.2e-2
        self.C_a = 4e2
        self.C_c = 5e2
        self.V_a = 5.5e-5
        self.V_c = 5.5e-5
        self.A_m = 5.0e-4
        self.Y_ac = 0.05
        self.K_dec = 8.33e-4
        self.f_x = 10.0
        self.alpha = 0.051
        self.beta = 0.063
        self.k1_0 = 0.207
        self.k2_0 = 3.288e-5
        self.K_AC = 0.592
        self.K_O2 = 0.004
        self.Q_a = 2.25e-5
        self.Q_c = 1.11e-3
        self.C_AC_in = 1.56
        self.C_CO2_in = 0.0
        self.X_in = 0.0
        self.C_H_in = 0.0
        self.C_O2_in = 0.3125
        self.C_M_in = 0.0
        self.C_OH_in = 0.0
        self.U0 = 0.77
        self.C_AC = 0.0
        self.C_CO2 = 0.0
        self.C_H = 0.0
        self.X = 0.0
        self.C_O2 = 0.0
        self.C_OH = 0.0
        self.C_M = 0.0
        self.eta_a = 0.0
        self.eta_c = 0.0
        
    fn __moveinit__(out self: Self, owned existing: Self):
        """
        Move initializer for MFCModel.
        This is called when an instance is moved, ensuring proper ownership transfer.
        """
        self.F = existing.F
        self.R = existing.R
        self.T = existing.T
        self.k_m = existing.k_m
        self.d_m = existing.d_m
        self.k_aq = existing.k_aq
        self.d_cell = existing.d_cell
        self.C_a = existing.C_a
        self.C_c = existing.C_c
        self.V_a = existing.V_a
        self.V_c = existing.V_c
        self.A_m = existing.A_m
        self.Y_ac = existing.Y_ac
        self.K_dec = existing.K_dec
        self.f_x = existing.f_x
        self.alpha = existing.alpha
        self.beta = existing.beta
        self.k1_0 = existing.k1_0
        self.k2_0 = existing.k2_0
        self.K_AC = existing.K_AC
        self.K_O2 = existing.K_O2
        self.Q_a = existing.Q_a        
        self.Q_c = existing.Q_c
        self.C_AC_in = existing.C_AC_in
        self.C_CO2_in = existing.C_CO2_in
        self.X_in = existing.X_in
        self.C_H_in = existing.C_H_in
        self.C_O2_in = existing.C_O2_in
        self.C_M_in = existing.C_M_in
        self.C_OH_in = existing.C_OH_in        
        self.U0 = existing.U0        
        self.C_AC = 0.0  # Reset to default value on move init.
        self.C_CO2 = 0.0  # Reset to default value on move init.
        self.C_H = 0.0  # Reset to default value on move init.
        self.X = 0.0  # Reset to default value on move init.
        self.C_O2 = 0.0  # Reset to default value on move init.
        self.C_OH = 0.0  # Reset to default value on move init.
        self.C_M = 0.0  # Reset to default value on move init.
        self.eta_a = 0.0  # Reset to default value on move init.
        self.eta_c = 0.0  # Reset to default value on move init.






    fn __copyinit__(out self: Self, existing: Self):
        """
        Copy initializer for MFCModel.
        This is called when an instance is copied, ensuring proper value transfer.
        """
        self.F = existing.F
        self.R = existing.R
        self.T = existing.T
        self.k_m = existing.k_m
        self.d_m = existing.d_m
        self.k_aq = existing.k_aq
        self.d_cell = existing.d_cell
        self.C_a = existing.C_a
        self.C_c = existing.C_c
        self.V_a = existing.V_a
        self.V_c = existing.V_c
        self.A_m = existing.A_m
        self.Y_ac = existing.Y_ac
        self.K_dec = existing.K_dec
        self.f_x = existing.f_x
        self.alpha = existing.alpha
        self.beta = existing.beta
        self.k1_0 = existing.k1_0
        self.k2_0 = existing.k2_0
        self.K_AC = existing.K_AC
        self.K_O2 = existing.K_O2
        self.Q_a = existing.Q_a        
        self.Q_c = existing.Q_c
        self.C_AC_in = existing.C_AC_in
        self.C_CO2_in = existing.C_CO2_in
        self.X_in = existing.X_in
        self.C_H_in = existing.C_H_in
        self.C_O2_in = existing.C_O2_in
        self.C_M_in = existing.C_M_in
        self.C_OH_in = existing.C_OH_in        
        self.U0 = existing.U0        
        self.C_AC = existing.C_AC
        self.C_CO2 = existing.C_CO2
        self.C_H = existing.C_H
        self.X = existing.X
        self.C_O2 = existing.C_O2
        self.C_OH = existing.C_OH
        self.C_M = existing.C_M
        self.eta_a = existing.eta_a
        self.eta_c = existing.eta_c

    fn __repr__(self: MFCModel) -> String:
        var repr = (
            "MFCModel("
            + "F={self.F}, R={self.R}, T={self.T}, k_m={self.k_m}, d_m={self.d_m}, "
            + "k_aq={self.k_aq}, d_cell={self.d_cell}, C_a={self.C_a}, C_c={self.C_c}, "
            + "V_a={self.V_a}, V_c={self.V_c}, A_m={self.A_m}, Y_ac={self.Y_ac}, "
            + "K_dec={self.K_dec}, f_x={self.f_x}, alpha={self.alpha}, beta={self.beta}, "
            + "k1_0={self.k1_0}, k2_0={self.k2_0}, K_AC={self.K_AC}, K_O2={self.K_O2}, "
            + "Q_a={self.Q_a}, Q_c={self.Q_c}, C_AC_in={self.C_AC_in}, C_CO2_in={self.C_CO2_in}, "
            + "X_in={self.X_in}, C_H_in={self.C_H_in}, C_O2_in={self.C_O2_in}, "
            + "C_M_in={self.C_M_in}, C_OH_in={self.C_OH_in}, U0={self.U0}, "
            + "C_AC={self.C_AC}, C_CO2={self.C_CO2}, C_H={self.C_H}, X={self.X}, "
            + "C_O2={self.C_O2}, C_OH={self.C_OH}, C_M={self.C_M}, "
            + "eta_a={self.eta_a}, eta_c={self.eta_c})"
        )
        return repr

    fn mfc_odes(self: MFCModel, t: Float64, y: List[Float64], i_fc: Float64) -> List[Float64]:
        """
        Calculates the derivatives. This now accepts a NumPy array directly
        and returns a Mojo List, which is automatically converted for Python.
        """
        # Add a check for the input array's length for robustness.
        try:
            if y._len != 9:
                # This error message will be propagated to Python as an exception.
                raise Error("Input list 'y' must have exactly 9 elements.")
        except:
            return List[Float64]()
    
        var C_AC = y[0]
        var C_CO2 = y[1]
        var C_H = y[2]
        var X = y[3]
        var C_O2 = y[4]
        var C_OH = y[5]
        var C_M = y[6]
        var eta_a = y[7]
        var eta_c = y[8]

        var r1 = self.k1_0 * math.exp((self.alpha * self.F) / (self.R * self.T) * eta_a) * (C_AC / (self.K_AC + C_AC)) * X 
        var r2 = -self.k2_0 * (C_O2 / (self.K_O2 + C_O2)) * math.exp((self.beta - 1.0) * self.F / (self.R * self.T) * eta_c)
        var N_M = (3600.0 * i_fc) / self.F

        var dC_AC_dt = (self.Q_a * (self.C_AC_in - C_AC) - self.A_m * r1) / self.V_a
        var dC_CO2_dt = (self.Q_a * (self.C_CO2_in - C_CO2) + 2.0 * self.A_m * r1) / self.V_a
        var dC_H_dt = (self.Q_a * (self.C_H_in - C_H) + 8.0 * self.A_m * r1) / self.V_a
        var dX_dt = (self.Q_a * (self.X_in - X) / self.f_x + self.A_m * self.Y_ac * r1) / self.V_a - self.K_dec * X

        var dC_O2_dt = (self.Q_c * (self.C_O2_in - C_O2) + r2 * self.A_m) / self.V_c
        var dC_OH_dt = (self.Q_c * (self.C_OH_in - C_OH) - 4.0 * r2 * self.A_m) / self.V_c
        var dC_M_dt = (self.Q_c * (self.C_M_in - C_M) + N_M * self.A_m) / self.V_c
        
        var d_eta_a_dt = (3600.0 * i_fc - 8.0 * self.F * r1) / self.C_a
        var d_eta_c_dt = (-3600.0 * i_fc - 4.0 * self.F * r2) / self.C_c

        var derivatives = List[Float64]()
        derivatives.append(dC_AC_dt)
        derivatives.append(dC_CO2_dt)
        derivatives.append(dC_H_dt)            
        derivatives.append(dX_dt)
        derivatives.append(dC_O2_dt)
        derivatives.append(dC_OH_dt)
        derivatives.append(dC_M_dt)
        derivatives.append(d_eta_a_dt)
        derivatives.append(d_eta_c_dt)
        return derivatives

and python file:

# With the July 2025 API, compiling the Mojo file creates a standard
# Python module. The legacy `max.mojo.importer` is no longer needed.


import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
import max.mojo.importer
import sys

sys.path.insert(0, "")
import odes  # Import the Mojo module compiled from odes.mojo
# 1. Initialize the Mojo model struct
mfc = odes.MFCModel()  # Create an instance of the MFCModel

# 2. Define initial conditions for the state variables [y0]
# C_AC, C_CO2, C_H, X, C_O2, C_OH, C_M, eta_a, eta_c
y0 = [
    1.5,      # C_AC: Start with near-influent concentration
    0.1,      # C_CO2: Small initial amount
    1e-4,     # C_H: Corresponds to pH 4
    0.1,      # X: Initial biomass
    0.3,      # C_O2: Start with near-influent concentration
    1e-7,     # C_OH: Corresponds to pH 7
    0.1,      # C_M: Small initial amount
    0.01,     # eta_a: Small initial anodic overpotential
    -0.01     # eta_c: Small initial cathodic overpotential
]

# 3. Define simulation time span
t_span = [0, 100]  # Simulate for 100 hours
t_eval = np.linspace(t_span[0], t_span[1], 500) # Points to evaluate solution

# 4. Set a constant current density for this example
constant_i_fc = 1.0 # A/m²

# 5. Run the ODE solver, passing the Mojo function as the "fun" argument.
#    The core calculations are now executed by the high-performance Mojo code.
solution = solve_ivp(
    fun=lambda t, y: mfc.mfc_odes(t, y, constant_i_fc),
    t_span=t_span,
    y0=y0,
    t_eval=t_eval,
    method='RK45'
)

# 6. Plot the results using Matplotlib
plt.style.use('seaborn-v0_8-whitegrid')
fig, ax = plt.subplots(figsize=(10, 6))
ax.plot(solution.t, solution.y[0], label='Acetate Concentration (C_AC)', color='b')
ax.set_xlabel('Time (hours)', fontsize=12)
ax.set_ylabel('Concentration (mol/m³)', fontsize=12)
ax.set_title('Mojo-Powered MFC Simulation: Acetate Concentration', fontsize=14)
ax.legend()
ax.grid(True)
plt.show()

System information

(mojo-practice) uge@uge-X670E-AORUS-XTREME:~/modular/mojo-practice/q-learning-mfcs$ pixi info
System

   Pixi version: 0.49.0
       Platform: linux-64

Virtual packages: __unix=0=0
: __linux=6.11.0=0
: __glibc=2.39=0
: __archspec=1=zen4
Cache dir: /home/uge/.cache/rattler/cache
Auth storage: /home/uge/.rattler/credentials.json
Config locations: No config files found

Global

        Bin dir: /home/uge/.pixi/bin
Environment dir: /home/uge/.pixi/envs
   Manifest dir: /home/uge/.pixi/manifests/pixi-global.toml

Workspace

           Name: mojo-practice
        Version: 0.1.0
  Manifest file: /home/uge/modular/mojo-practice/pixi.toml
   Last updated: 08-07-2025 00:41:49

Environments

    Environment: default
       Features: default
       Channels: conda-forge, https://conda.modular.com/max-nightly, https://repo.prefix.dev/modular, https://repo.prefix.dev/mojo, https://repo.prefix.dev/modular-community

Dependency count: 6
Dependencies: scipy, matplotlib, python, numpy, modular, magic
Target platforms: linux-aarch64, linux-64
Prefix location: /home/uge/modular/mojo-practice/.pixi/envs/default

Package Version Build Size Kind Source
max-core 25.5.0.dev2025070705 release 214 MiB conda https://conda.modular.com/max-nightly/
max-pipelines 25.5.0.dev2025070705 release 9.8 KiB conda https://conda.modular.com/max-nightly/
max-python 25.5.0.dev2025070705 release 30.2 MiB conda https://conda.modular.com/max-nightly/

Start with the provided Mojo example at @export PyInit_odes, the PythonModuleBuilder setup, and the Python call to mfc.MFCModel.mfc_odes; reproduce the missing-method behavior with the shown Python and Mojo entry points. Done means the compiled module exposes the MFCModel methods and accepts the demonstrated return type and arguments without the current PythonObject restriction.