From 987b92ab1e2bd3778b754217cd3208143d28bbbe Mon Sep 17 00:00:00 2001
From: Koen Vlaswinkel <koesie10@github.com>
Date: Tue, 2 Jul 2024 16:18:22 +0200
Subject: [PATCH] Enable type models for Python

---
 .../model-editor/languages/python/index.ts    | 25 +++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/extensions/ql-vscode/src/model-editor/languages/python/index.ts b/extensions/ql-vscode/src/model-editor/languages/python/index.ts
index 691687f1b..d6611d4ea 100644
--- a/extensions/ql-vscode/src/model-editor/languages/python/index.ts
+++ b/extensions/ql-vscode/src/model-editor/languages/python/index.ts
@@ -173,6 +173,31 @@ export const python: ModelsAsDataLanguage = {
         };
       },
     },
+    type: {
+      extensiblePredicate: "typeModel",
+      // extensible predicate typeModel(string type1, string type2, string path);
+      generateMethodDefinition: (method) => [
+        method.relatedTypeName,
+        pythonType(method.packageName, method.typeName, method.endpointType),
+        pythonPath(method.methodName, method.path),
+      ],
+      readModeledMethod: (row) => {
+        const { packageName, typeName, methodName, endpointType, path } =
+          parsePythonTypeAndPath(row[1] as string, row[2] as string);
+
+        return {
+          type: "type",
+          relatedTypeName: row[0] as string,
+          path,
+          signature: pythonMethodSignature(typeName, methodName),
+          endpointType,
+          packageName,
+          typeName,
+          methodName,
+          methodParameters: "",
+        };
+      },
+    },
   },
   getArgumentOptions: (method) => {
     // Argument and Parameter are equivalent in Python, but we'll use Argument in the model editor