feat: wire up python schema extraction (#3369)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Justfile

@@ -29,6 +29,7 @@ clean:
   rm -rf node_modules
   rm -rf frontend/console/dist
   rm -rf frontend/console/node_modules
+  rm -rf python-runtime/ftl/.venv
   find . -name '*.zip' -exec rm {} \;
   mvn -f jvm-runtime/ftl-runtime clean
 

python-runtime/compile/build-template/main/main.py.tmpl

@@ -1,3 +1,3 @@
-// Code generated by FTL. DO NOT EDIT.
+# Code generated by FTL. DO NOT EDIT.
 
 {{- $name := .Name -}}

python-runtime/compile/build.go

@@ -3,16 +3,20 @@ package compile
 import (
 	"context"
 	"fmt"
+	"os"
 	"path"
 	"path/filepath"
 	stdreflect "reflect"
 	"strconv"
 	"strings"
 
 	"github.com/block/scaffolder"
+	"google.golang.org/protobuf/proto"
 
+	schemapb "github.com/TBD54566975/ftl/backend/protos/xyz/block/ftl/v1/schema"
 	"github.com/TBD54566975/ftl/internal"
 	"github.com/TBD54566975/ftl/internal/builderrors"
+	"github.com/TBD54566975/ftl/internal/exec"
 	"github.com/TBD54566975/ftl/internal/log"
 	"github.com/TBD54566975/ftl/internal/moduleconfig"
 	"github.com/TBD54566975/ftl/internal/schema"
@@ -40,14 +44,33 @@ func Build(ctx context.Context, projectRootDir, stubsRoot string, config modulec
 
 	buildDir := buildDir(config.Dir)
 
-	// TODO: call the python schema extractor. grab the output of le script. unmarshal into schema proto. unmarshal that into go type. return
-	// same with build errors
+	// Execute the Python schema extractor
+	if err := exec.Command(ctx, log.Debug, config.Dir, "uv", "run", "-m", "ftl.cli.schema_extractor", ".").RunBuffered(ctx); err != nil {
+		return nil, nil, fmt.Errorf("failed to extract schema: %w", err)
+	}
+
+	outputFile := filepath.Join(buildDir, "schema.pb")
+	serializedData, err := os.ReadFile(outputFile)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to read serialized schema: %w", err)
+	}
+
+	var modulepb schemapb.Module
+	err = proto.Unmarshal(serializedData, &modulepb)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to unmarshal module proto: %w", err)
+	}
+
+	module, err := schema.ModuleFromProto(&modulepb)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to deserialize module schema: %w", err)
+	}
 
 	if err := internal.ScaffoldZip(buildTemplateFiles(), buildDir, mctx, scaffolder.Functions(scaffoldFuncs)); err != nil {
 		return moduleSch, nil, fmt.Errorf("failed to scaffold build template: %w", err)
 	}
 
-	return nil, nil, nil
+	return module, nil, nil
 }
 
 var scaffoldFuncs = scaffolder.FuncMap{

python-runtime/compile/build_test.go

@@ -0,0 +1,55 @@
+package compile
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/TBD54566975/ftl/internal/exec"
+	"github.com/TBD54566975/ftl/internal/log"
+	"github.com/TBD54566975/ftl/internal/moduleconfig"
+	"github.com/TBD54566975/ftl/internal/schema"
+	"github.com/alecthomas/assert/v2"
+)
+
+func TestBuild(t *testing.T) {
+	// set up
+	moduleDir, err := filepath.Abs("testdata/echo")
+	assert.NoError(t, err)
+	ctx := log.ContextWithNewDefaultLogger(context.Background())
+	assert.NoError(t, os.RemoveAll(filepath.Join(moduleDir, ".venv")))
+	assert.NoError(t, exec.Command(ctx, log.Debug, moduleDir, "uv", "sync", "-n").Run())
+
+	t.Run("schema extraction", func(t *testing.T) {
+		config := moduleconfig.AbsModuleConfig{
+			Dir:    moduleDir,
+			Module: "test",
+		}
+		actual, buildErrors, err := Build(ctx, "", "", config, &schema.Schema{}, nil, true)
+		assert.NoError(t, err)
+		assert.Equal(t, 0, len(buildErrors))
+		expected := &schema.Module{
+			Name: "echo",
+			Decls: []schema.Decl{
+				&schema.Data{
+					Name: "EchoRequest",
+					Fields: []*schema.Field{
+						{Name: "name", Type: &schema.String{}},
+					},
+				},
+				&schema.Data{
+					Name: "EchoResponse",
+					Fields: []*schema.Field{
+						{Name: "message", Type: &schema.String{}},
+					},
+				},
+				&schema.Verb{Name: "echo",
+					Request:  &schema.Ref{Module: "echo", Name: "EchoRequest"},
+					Response: &schema.Ref{Module: "echo", Name: "EchoResponse"},
+				},
+			},
+		}
+		assert.Equal(t, expected, actual, assert.Exclude[schema.Position]())
+	})
+}

python-runtime/compile/testdata/echo/.python-version

@@ -0,0 +1 @@
+3.12

python-runtime/compile/testdata/echo/echo.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass
+
+from ftl import verb
+
+
+@dataclass
+class EchoRequest:
+    name: str
+
+
+@dataclass
+class EchoResponse:
+    message: str
+
+
+@verb
+def echo(req: EchoRequest) -> EchoResponse:
+    return EchoResponse(message=f"ayooo, {req.name}!")

python-runtime/compile/testdata/echo/ftl.toml

@@ -0,0 +1,2 @@
+module = "echo"
+language = "python"

python-runtime/compile/testdata/echo/pyproject.toml

@@ -0,0 +1,10 @@
+[project]
+name = "echo"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = ["ftl"]
+
+[tool.uv.sources]
+ftl = { path = "../../../../python-runtime/ftl" }

python-runtime/ftl/src/ftl/cli/schema_extractor.py

@@ -3,6 +3,7 @@
 import concurrent.futures
 import os
 import sys
+import tomllib
 from contextlib import contextmanager
 
 from ftl.extract import (
@@ -13,15 +14,14 @@
     VerbExtractor,
 )
 
-# analyzers is now a list of lists, where each sublist contains analyzers that can run in parallel
+# analyzers is a list of lists, where each sublist contains analyzers that can run in parallel
 analyzers = [
     [VerbExtractor],
     [TransitiveExtractor],
 ]
 
-
 @contextmanager
-def set_analysis_mode(path):
+def set_analysis_mode(path: str):
     original_sys_path = sys.path.copy()
     sys.path.append(path)
     try:
@@ -30,7 +30,22 @@ def set_analysis_mode(path):
         sys.path = original_sys_path
 
 
-def analyze_directory(module_dir):
+def get_module_name(ftl_dir: str) -> str:
+    ftl_toml_path = os.path.join(ftl_dir, "ftl.toml")
+
+    if not os.path.isfile(ftl_toml_path):
+        raise FileNotFoundError(f"ftl.toml file not found in the specified module directory: {ftl_dir}")
+
+    with open(ftl_toml_path, "rb") as f:
+        config = tomllib.load(f)
+        module_name = config.get("module")
+        if module_name:
+            return module_name
+        else:
+            raise ValueError("module name not found in ftl.toml")
+
+
+def analyze_directory(module_dir: str):
     """Analyze all Python files in the given module_dir in parallel."""
     global_ctx = GlobalExtractionContext()
 
@@ -55,14 +70,17 @@ def analyze_directory(module_dir):
                     future.result()  # raise any exception that occurred in the worker process
                 except Exception as exc:
                     print(f"failed to extract schema from {file_path}: {exc};")
-                # else:
-                #     print(f"File {file_path} analyzed successfully.")
 
-    for ref_key, decl in global_ctx.deserialize().items():
-        print(f"Extracted Decl:\n{decl}")
+    output_dir = os.path.join(module_dir, ".ftl")
+    os.makedirs(output_dir, exist_ok=True)  # Create .ftl directory if it doesn't exist
+    output_file = os.path.join(output_dir, "schema.pb")
+
+    serialized_schema = global_ctx.to_module_schema(get_module_name(module_dir)).SerializeToString()
+    with open(output_file, "wb") as f:
+        f.write(serialized_schema)
 
 
-def analyze_file(global_ctx: GlobalExtractionContext, file_path, analyzer_batch):
+def analyze_file(global_ctx: GlobalExtractionContext, file_path: str, analyzer_batch):
     """Analyze a single Python file using multiple analyzers in parallel."""
     module_name = os.path.splitext(os.path.basename(file_path))[0]
     file_ast = ast.parse(open(file_path).read())

python-runtime/ftl/src/ftl/extract/__init__.py

@@ -1,13 +1,13 @@
-from .common import (
+from ftl.extract.common import (
     extract_basic_type,
     extract_class_type,
     extract_function_type,
     extract_map,
     extract_slice,
     extract_type,
 )
-from .context import GlobalExtractionContext, LocalExtractionContext
-from .transitive import TransitiveExtractor
+from ftl.extract.context import GlobalExtractionContext, LocalExtractionContext
+from ftl.extract.transitive import TransitiveExtractor
 
 __all__ = [
     "extract_type",

python-runtime/ftl/src/ftl/extract/common.py

@@ -2,7 +2,7 @@
 
 from ftl.protos.xyz.block.ftl.v1.schema import schema_pb2 as schemapb
 
-from .context import LocalExtractionContext
+from ftl.extract.context import LocalExtractionContext
 
 
 def extract_type(
@@ -86,14 +86,18 @@ def extract_basic_type(type_hint: Type[Any]) -> Optional[schemapb.Type]:
 def extract_class_type(
     local_ctx: LocalExtractionContext, type_hint: Type[Any]
 ) -> Optional[schemapb.Type]:
-    ref = schemapb.Ref(name=type_hint.__name__, module=type_hint.__module__)
+    ref = schemapb.Ref(name=get_base_module_name(type_hint.__name__), module=type_hint.__module__)
     local_ctx.add_needs_extraction(ref)
     return schemapb.Type(ref=ref)
 
 
 def extract_function_type(
     local_ctx: LocalExtractionContext, type_hint: Type[Any]
 ) -> Optional[schemapb.Type]:
-    ref = schemapb.Ref(name=type_hint.__name__, module=type_hint.__module__)
+    ref = schemapb.Ref(name=get_base_module_name(type_hint.__name__), module=type_hint.__module__)
     local_ctx.add_needs_extraction(ref)
     return schemapb.Type(ref=ref)
+
+def get_base_module_name(fq_module_name: str) -> str:
+    """Return the base (root) module name from a fully qualified module name."""
+    return fq_module_name.split('.')[0]