add user defined functions to power-puter

py/power_puter.py

@@ -445,6 +445,119 @@ def _get_input_node(self, input_name, node=None):
       log_node_warn(_NODE_NAME, f'No input node found for "{input_name}". ')
     return None
 
+  # --- DRY helpers for function/lambda signature handling ---
+  def _prepare_signature(self, args_node, outer_ctx):
+    """Prepare a callable signature from an ast.arguments node.
+
+    Returns a dict with:
+      param_names, defaults_map, kwonly_names, kw_defaults_map, vararg_name, kwarg_name
+    """
+    posonly = [a.arg for a in getattr(args_node, 'posonlyargs', [])]
+    normal = [a.arg for a in args_node.args]
+    param_names = posonly + normal
+
+    # Positional defaults map to last N positional params
+    defaults_evaluated = []
+    if getattr(args_node, 'defaults', None):
+      for d in args_node.defaults:
+        defaults_evaluated.append(self._eval_statement(d, ctx=outer_ctx))
+    defaults_map = {}
+    if defaults_evaluated:
+      for i, val in enumerate(defaults_evaluated):
+        param = param_names[len(param_names) - len(defaults_evaluated) + i]
+        defaults_map[param] = val
+
+    kwonly_names = [a.arg for a in getattr(args_node, 'kwonlyargs', [])]
+    kw_defaults_evaluated = []
+    if getattr(args_node, 'kw_defaults', None):
+      for d in args_node.kw_defaults:
+        if d is None:
+          kw_defaults_evaluated.append(None)
+        else:
+          kw_defaults_evaluated.append(self._eval_statement(d, ctx=outer_ctx))
+    kw_defaults_map = {}
+    if kwonly_names:
+      for i, name in enumerate(kwonly_names):
+        if i < len(kw_defaults_evaluated):
+          kw_defaults_map[name] = kw_defaults_evaluated[i]
+        else:
+          kw_defaults_map[name] = None
+
+    vararg_name = args_node.vararg.arg if args_node.vararg else None
+    kwarg_name = args_node.kwarg.arg if args_node.kwarg else None
+
+    return {
+      'param_names': param_names,
+      'defaults_map': defaults_map,
+      'kwonly_names': kwonly_names,
+      'kw_defaults_map': kw_defaults_map,
+      'vararg_name': vararg_name,
+      'kwarg_name': kwarg_name,
+    }
+
+  def _bind_call(self, signature: dict, call_args, call_kwargs):
+    """Bind provided arguments to a prepared signature, enforcing the same rules
+    as the inlined logic previously used for FunctionDef and Lambda.
+    Returns a dict of bound parameters (including varargs/kwargs names if present).
+    """
+    param_names = signature['param_names']
+    defaults_map = signature['defaults_map']
+    kwonly_names = signature['kwonly_names']
+    kw_defaults_map = signature['kw_defaults_map']
+    vararg_name = signature['vararg_name']
+    kwarg_name = signature['kwarg_name']
+
+    bindings = {}
+    work_kwargs = dict(call_kwargs)
+
+    # Positional parameters
+    positionally_bound = set()
+    for i, pname in enumerate(param_names):
+      if i < len(call_args):
+        bindings[pname] = call_args[i]
+        positionally_bound.add(pname)
+      else:
+        if pname in defaults_map:
+          bindings[pname] = defaults_map[pname]
+        else:
+          raise TypeError(f"Missing required positional argument: {pname}")
+
+    # varargs
+    if vararg_name:
+      bindings[vararg_name] = tuple(call_args[len(param_names):])
+    else:
+      if len(call_args) > len(param_names):
+        raise TypeError("Too many positional arguments")
+
+    # Keyword-only args
+    for kname in kwonly_names:
+      if kname in work_kwargs:
+        bindings[kname] = work_kwargs.pop(kname)
+      elif kname in kw_defaults_map and kw_defaults_map[kname] is not None:
+        bindings[kname] = kw_defaults_map[kname]
+      else:
+        raise TypeError(f"Missing required keyword-only argument: {kname}")
+
+    # Remaining kwargs and positional-by-name
+    if kwarg_name:
+      for pname in param_names:
+        if pname in work_kwargs:
+          if pname in positionally_bound:
+            raise TypeError(f"Multiple values for argument: {pname}")
+          bindings[pname] = work_kwargs.pop(pname)
+      bindings[kwarg_name] = {**work_kwargs}
+    else:
+      for pname in param_names:
+        if pname in work_kwargs:
+          if pname in positionally_bound:
+            raise TypeError(f"Multiple values for argument: {pname}")
+          bindings[pname] = work_kwargs.pop(pname)
+      if work_kwargs:
+        unknown = ', '.join(work_kwargs.keys())
+        raise TypeError(f"Unexpected keyword arguments: {unknown}")
+
+    return bindings
+
   def _eval_statement(self, stmt: ast.AST, ctx: dict, prev_stmt: Union[ast.AST, None] = None):
     """Evaluates an ast.stmt."""
 
@@ -459,7 +572,8 @@ def _eval_statement(self, stmt: ast.AST, ctx: dict, prev_stmt: Union[ast.AST, No
       return self._eval_statement(stmt.value, ctx=ctx)
 
     if isinstance(stmt, (ast.Constant, ast.Num)):
-      return stmt.n
+      # ast.Constant (Py>=3.8) uses .value; ast.Num (Py<3.8/compat) uses .n
+      return getattr(stmt, 'value', getattr(stmt, 'n', None))
 
     if isinstance(stmt, ast.BinOp):
       left = self._eval_statement(stmt.left, ctx=ctx)
@@ -650,6 +764,21 @@ def handle_gen(generators: list[ast.comprehension]):
       handle_gen(generators)
       return final_list
 
+    # Lambda expression support: build a callable that captures current context
+    if isinstance(stmt, ast.Lambda):
+      lam = stmt
+      outer_ctx = ctx
+
+      signature = self._prepare_signature(lam.args, outer_ctx)
+
+      def _lambda_fn(*call_args, **call_kwargs):
+        lctx = {**outer_ctx}
+        bindings = self._bind_call(signature, call_args, call_kwargs)
+        lctx.update(bindings)
+        return self._eval_statement(lam.body, ctx=lctx)
+
+      return _lambda_fn
+
     if isinstance(stmt, ast.Call):
       call = None
       args = []
@@ -667,6 +796,9 @@ def handle_gen(generators: list[ast.comprehension]):
         name = stmt.func.id
         if name in _BUILT_INS:
           call = _BUILT_INS[name]
+        else:
+          # Evaluate the name from the current context to allow calling user-defined functions
+          call = self._eval_statement(stmt.func, ctx=ctx)
 
       if isinstance(call, str) and call.startswith(_BUILTIN_FN_PREFIX):
         fn = _get_built_in_fn_by_key(call)
@@ -681,10 +813,14 @@ def handle_gen(generators: list[ast.comprehension]):
       if not call:
         raise ValueError(f'No call for ast.Call {name}')
 
+      if not callable(call):
+        raise TypeError(f'Attempted to call a non-callable object: {call}')
+
       for arg in stmt.args:
         args.append(self._eval_statement(arg, ctx=ctx))
       for kwarg in stmt.keywords:
         kwargs[kwarg.arg] = self._eval_statement(kwarg.value, ctx=ctx)
+      # noinspection PyCallingNonCallable
       return call(*args, **kwargs)
 
     if isinstance(stmt, ast.Compare):
@@ -752,6 +888,44 @@ def handle_gen(generators: list[ast.comprehension]):
       ctx[stmt.target.id] = value
       return value
 
+    if isinstance(stmt, ast.FunctionDef):
+      # Define a user function: store a callable in the current context under the function name.
+      fn_def = stmt
+
+      # Capture outer context by reference for reads; we'll use a new local ctx per invocation.
+      outer_ctx = ctx
+
+      signature = self._prepare_signature(fn_def.args, outer_ctx)
+
+      def _user_fn(*call_args, **call_kwargs):
+        # Local context starts as a shallow copy of the outer context
+        lctx = {**outer_ctx}
+        # Ensure we don't inherit a return marker from the defining context
+        if '__returned__' in lctx:
+          del lctx['__returned__']
+
+        bindings = self._bind_call(signature, call_args, call_kwargs)
+
+        # Update local context with bound params and function name (for recursion)
+        lctx.update(bindings)
+        lctx[fn_def.name] = _user_fn
+
+        # Execute function body
+        ret_val = None
+        for b in fn_def.body:
+          ret_val = self._eval_statement(b, ctx=lctx)
+          if '__returned__' in lctx:
+            ret_val = lctx['__returned__']
+            break
+        # Clean return marker if present
+        if '__returned__' in lctx:
+          del lctx['__returned__']
+        return ret_val
+
+      # Store function in current context
+      ctx[fn_def.name] = _user_fn
+      return None
+
     if isinstance(stmt, ast.Return):
       if stmt.value is None:
         value = None