Repository: MotiaDev/motia
File: packages/core/src/python/python-runner.py (Python)

Code before:
import json
import importlib.util
import traceback
import asyncio
import os
from logger import Logger
from rpc import RpcSender
from rpc_state_manager import RpcStateManager

from typing import Any

def parse_args(arg: str) -> Any:
    from types import SimpleNamespace
    
    try:
        return json.loads(arg, object_hook=lambda d: SimpleNamespace(**d))
    except json.JSONDecodeError:
        print('Error parsing args:', arg)
        return arg

class Context:
    def __init__(self, args: Any, sender: RpcSender, file_name: str):
        self.trace_id = args.traceId
        self.flows = args.flows
        self.file_name = file_name
        self.sender = sender
        self.state = RpcStateManager(self.sender)
        self.logger = Logger(self.trace_id, self.flows, self.file_name, self.sender)

    async def emit(self, event: Any):
        return await self.sender.send('emit', event)

async def run_python_module(file_path: str, rpc: RpcSender, args: Any) -> None:
    try:
        # Construct path relative to steps directory
        flows_dir = os.path.join(os.getcwd(), 'steps')
        module_path = os.path.join(flows_dir, file_path)

        # Load the module dynamically
        spec = importlib.util.spec_from_file_location("dynamic_module", module_path)
        if spec is None or spec.loader is None:
            raise ImportError(f"Could not load module from {module_path}")
            
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)

        # Check if the handler function exists
        if not hasattr(module, 'handler'):
            raise AttributeError(f"Function 'handler' not found in module {module_path}")

        context = Context(args, rpc, file_path)

Review comment (suggestion):
nitpick: `sender` doesn't sound very intuitive, I would suggest changing this for state manager or soemthing more intuitive. AFAIK the rpc implementation replaced the http state manager, so maybe renaming it to `state` or `stateManager` might be best

Code after:
import json
import importlib.util
import traceback
import asyncio
import os
from logger import Logger
from rpc import RpcSender
from rpc_state_manager import RpcStateManager

from typing import Any

def parse_args(arg: str) -> Any:
    from types import SimpleNamespace
    
    try:
        return json.loads(arg, object_hook=lambda d: SimpleNamespace(**d))
    except json.JSONDecodeError:
        print('Error parsing args:', arg)
        return arg

class Context:
    def __init__(self, args: Any, rpc: RpcSender, file_name: str):
        self.trace_id = args.traceId
        self.flows = args.flows
        self.file_name = file_name
        self.rpc = rpc
        self.state = RpcStateManager(rpc)
        self.logger = Logger(self.trace_id, self.flows, self.file_name, rpc)

    async def emit(self, event: Any):
        return await self.rpc.send('emit', event)

async def run_python_module(file_path: str, rpc: RpcSender, args: Any) -> None:
    try:
        # Construct path relative to steps directory
        flows_dir = os.path.join(os.getcwd(), 'steps')
        module_path = os.path.join(flows_dir, file_path)

        # Load the module dynamically
        spec = importlib.util.spec_from_file_location("dynamic_module", module_path)
        if spec is None or spec.loader is None:
            raise ImportError(f"Could not load module from {module_path}")
            
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)

        # Check if the handler function exists
        if not hasattr(module, 'handler'):
            raise AttributeError(f"Function 'handler' not found in module {module_path}")

        context = Context(args, rpc, file_path)