fix: typing of pdl_context in notebook extension and ast (#974)

Signed-off-by: Louis Mandel <[email protected]>

Author: mandel

pdl-live-react/src/pdl_ast.d.ts

@@ -2406,16 +2406,6 @@ export type PdlModelInput =
  *
  */
 export type Platform = "granite-io"
-/**
- * Backend name and configuration.
- *
- */
-export type Backend =
-  | LocalizedExpression
-  | string
-  | {
-      [k: string]: unknown
-    }
 /**
  * Parameters sent to the model.
  *
@@ -3220,7 +3210,7 @@ export interface GraniteioModelBlock {
   pdl__usage?: PdlUsage | null
   pdl__model_input?: PdlModelInput
   platform?: Platform
-  backend: Backend
+  backend: unknown
   processor?: unknown
   parameters?: Parameters
 }

pdl-live-react/src/pdl_ast_utils.ts

@@ -1,6 +1,6 @@
 import { match, P } from "ts-pattern"
 
-import { Backend, PdlBlock } from "./pdl_ast"
+import { PdlBlock } from "./pdl_ast"
 import { ExpressionT, isArgs } from "./helpers"
 
 export function map_block_children(
@@ -72,8 +72,7 @@ export function map_block_children(
         const parameters: Parameters = block.parameters
           ? f_expr(block.parameters)
           : undefined
-        // @ts-expect-error: f_expr does not preserve the type of the expression
-        const backend: Backend = f_expr(block.backend)
+        const backend = f_expr(block.backend)
         const processor = block.processor ? f_expr(block.processor) : undefined
         return {
           ...block,

src/pdl/pdl-schema.json

@@ -4612,7 +4612,8 @@
             {
               "additionalProperties": true,
               "type": "object"
-            }
+            },
+            {}
           ],
           "description": "Backend name and configuration.\n    ",
           "title": "Backend"

src/pdl/pdl_ast.py

@@ -551,7 +551,7 @@ class GraniteioModelBlock(ModelBlock):
     model: ExpressionType[str]
     """Model name used by the backend.
     """
-    backend: ExpressionType[str | dict[str, Any | object]]
+    backend: ExpressionType[str | dict[str, Any] | object]
     """Backend name and configuration.
     """
     processor: Optional[ExpressionType[str | object]] = None

src/pdl/pdl_notebook_ext.py

@@ -7,8 +7,9 @@
 
 from .pdl import InterpreterConfig, exec_str
 from .pdl_ast import get_default_model_parameters
+from .pdl_context import DependentContext
 from .pdl_dumper import block_to_dict
-from .pdl_lazy import PdlDict, PdlList
+from .pdl_lazy import PdlDict
 
 
 @magics_class
@@ -33,10 +34,12 @@ def pdl(self, line, cell, local_ns):
         line = line.strip()
         args = parse_argstring(self.pdl, line)
         if args.reset_context:
-            scope = local_ns | {"pdl_context": PdlList([])}
+            scope = local_ns | {"pdl_context": DependentContext([])}
         else:
             # local_ns won't be lazy; make it lazy again
-            scope = local_ns | {"pdl_context": PdlList(local_ns.get("pdl_context", []))}
+            scope = local_ns | {
+                "pdl_context": local_ns.get("pdl_context", DependentContext([]))
+            }
 
         if "pdl_model_default_parameters" not in scope:
             scope["pdl_model_default_parameters"] = get_default_model_parameters()