Use rich type columns and endpoint kinds for Python model editor

extensions/ql-vscode/src/model-editor/bqrs.ts

@@ -33,6 +33,7 @@ export function decodeBqrsToMethods(
     let libraryVersion: string | undefined;
     let type: ModeledMethodType;
     let classification: CallClassification;
+    let endpointKindColumn: string | BqrsEntityValue | undefined;
     let endpointType: EndpointType | undefined = undefined;
 
     if (mode === Mode.Application) {
@@ -47,6 +48,7 @@ export function decodeBqrsToMethods(
         libraryVersion,
         type,
         classification,
+        endpointKindColumn,
       ] = tuple as ApplicationModeTuple;
     } else {
       [
@@ -58,6 +60,7 @@ export function decodeBqrsToMethods(
         supported,
         library,
         type,
+        endpointKindColumn,
       ] = tuple as FrameworkModeTuple;
 
       classification = CallClassification.Unknown;
@@ -68,13 +71,18 @@ export function decodeBqrsToMethods(
     }
 
     if (definition.endpointTypeForEndpoint) {
-      endpointType = definition.endpointTypeForEndpoint({
-        endpointType,
-        packageName,
-        typeName,
-        methodName,
-        methodParameters,
-      });
+      endpointType = definition.endpointTypeForEndpoint(
+        {
+          endpointType,
+          packageName,
+          typeName,
+          methodName,
+          methodParameters,
+        },
+        typeof endpointKindColumn === "object"
+          ? endpointKindColumn.label
+          : endpointKindColumn,
+      );
     }
 
     if (endpointType === undefined) {

extensions/ql-vscode/src/model-editor/languages/models-as-data.ts

@@ -174,11 +174,14 @@ export type ModelsAsDataLanguage = {
    * be determined by heuristics.
    * @param method The method to get the endpoint type for. The endpoint type can be undefined if the
    *               query does not return an endpoint type.
+   * @param endpointKind An optional column that may be provided by the query to help determine the
+   *                     endpoint type.
    */
   endpointTypeForEndpoint?: (
     method: Omit<MethodDefinition, "endpointType"> & {
       endpointType: EndpointType | undefined;
     },
+    endpointKind: string | undefined,
   ) => EndpointType | undefined;
   predicates: ModelsAsDataLanguagePredicates;
   modelGeneration?: ModelsAsDataLanguageModelGeneration;

extensions/ql-vscode/src/model-editor/languages/python/access-paths.ts

@@ -4,7 +4,26 @@ import { EndpointType } from "../../method";
 
 const memberTokenRegex = /^Member\[(.+)]$/;
 
-export function parsePythonAccessPath(path: string): {
+// In Python, the type can contain both the package name and the type name.
+export function parsePythonType(type: string) {
+  // The first part is always the package name. All remaining parts are the type
+  // name.
+
+  const parts = type.split(".");
+  const packageName = parts.shift() ?? "";
+
+  return {
+    packageName,
+    typeName: parts.join("."),
+  };
+}
+
+// The type name can also be specified in the type, so this will combine
+// the already parsed type name and the type name from the access path.
+export function parsePythonAccessPath(
+  path: string,
+  shortTypeName: string,
+): {
   typeName: string;
   methodName: string;
   endpointType: EndpointType;
@@ -13,8 +32,12 @@ export function parsePythonAccessPath(path: string): {
   const tokens = parseAccessPathTokens(path);
 
   if (tokens.length === 0) {
+    const typeName = shortTypeName.endsWith("!")
+      ? shortTypeName.slice(0, -1)
+      : shortTypeName;
+
     return {
-      typeName: "",
+      typeName,
       methodName: "",
       endpointType: EndpointType.Method,
       path: "",
@@ -23,6 +46,10 @@ export function parsePythonAccessPath(path: string): {
 
   const typeParts = [];
   let endpointType = EndpointType.Function;
+  // If a short type name was given and it doesn't end in a `!`, then this refers to a method.
+  if (shortTypeName !== "" && !shortTypeName.endsWith("!")) {
+    endpointType = EndpointType.Method;
+  }
 
   let remainingTokens: typeof tokens = [];
 
@@ -32,6 +59,7 @@ export function parsePythonAccessPath(path: string): {
     if (memberMatch) {
       typeParts.push(memberMatch[1]);
     } else if (token.text === "Instance") {
+      // Alternative way of specifying that this refers to a method.
       endpointType = EndpointType.Method;
     } else {
       remainingTokens = tokens.slice(i);
@@ -40,9 +68,22 @@ export function parsePythonAccessPath(path: string): {
   }
 
   const methodName = typeParts.pop() ?? "";
-  const typeName = typeParts.join(".");
+  let typeName = typeParts.join(".");
   const remainingPath = remainingTokens.map((token) => token.text).join(".");
 
+  if (shortTypeName !== "") {
+    if (shortTypeName.endsWith("!")) {
+      // The actual type name is the name without the `!`.
+      shortTypeName = shortTypeName.slice(0, -1);
+    }
+
+    if (typeName !== "") {
+      typeName = `${shortTypeName}.${typeName}`;
+    } else {
+      typeName = shortTypeName;
+    }
+  }
+
   return {
     methodName,
     typeName,
@@ -51,53 +92,59 @@ export function parsePythonAccessPath(path: string): {
   };
 }
 
-export function pythonMethodSignature(typeName: string, methodName: string) {
-  return `${typeName}#${methodName}`;
-}
+export function parsePythonTypeAndPath(
+  type: string,
+  path: string,
+): {
+  packageName: string;
+  typeName: string;
+  methodName: string;
+  endpointType: EndpointType;
+  path: string;
+} {
+  const { packageName, typeName: shortTypeName } = parsePythonType(type);
+  const {
+    typeName,
+    methodName,
+    endpointType,
+    path: remainingPath,
+  } = parsePythonAccessPath(path, shortTypeName);
 
-function pythonTypePath(typeName: string) {
-  if (typeName === "") {
-    return "";
-  }
+  return {
+    packageName,
+    typeName,
+    methodName,
+    endpointType,
+    path: remainingPath,
+  };
+}
 
-  return typeName
-    .split(".")
-    .map((part) => `Member[${part}]`)
-    .join(".");
+export function pythonMethodSignature(typeName: string, methodName: string) {
+  return `${typeName}#${methodName}`;
 }
 
-export function pythonMethodPath(
+export function pythonType(
+  packageName: string,
   typeName: string,
-  methodName: string,
   endpointType: EndpointType,
 ) {
-  if (methodName === "") {
-    return pythonTypePath(typeName);
+  if (typeName !== "" && packageName !== "") {
+    return `${packageName}.${typeName}${endpointType === EndpointType.Function ? "!" : ""}`;
   }
 
-  const typePath = pythonTypePath(typeName);
-
-  let result = typePath;
-  if (typePath !== "" && endpointType === EndpointType.Method) {
-    result += ".Instance";
-  }
+  return `${packageName}${typeName}`;
+}
 
-  if (result !== "") {
-    result += ".";
+export function pythonMethodPath(methodName: string) {
+  if (methodName === "") {
+    return "";
   }
 
-  result += `Member[${methodName}]`;
-
-  return result;
+  return `Member[${methodName}]`;
 }
 
-export function pythonPath(
-  typeName: string,
-  methodName: string,
-  endpointType: EndpointType,
-  path: string,
-) {
-  const methodPath = pythonMethodPath(typeName, methodName, endpointType);
+export function pythonPath(methodName: string, path: string) {
+  const methodPath = pythonMethodPath(methodName);
   if (methodPath === "") {
     return path;
   }
@@ -111,7 +158,24 @@ export function pythonPath(
 
 export function pythonEndpointType(
   method: Omit<MethodDefinition, "endpointType">,
+  endpointKind: string | undefined,
 ): EndpointType {
+  switch (endpointKind) {
+    case "Function":
+      return EndpointType.Function;
+    case "InstanceMethod":
+      return EndpointType.Method;
+    case "ClassMethod":
+      return EndpointType.ClassMethod;
+    case "StaticMethod":
+      return EndpointType.StaticMethod;
+    case "InitMethod":
+      return EndpointType.Constructor;
+    case "Class":
+      return EndpointType.Class;
+  }
+
+  // Legacy behavior for when the kind column is missing.
   if (
     method.methodParameters.startsWith("(self,") ||
     method.methodParameters === "(self)"
@@ -120,3 +184,12 @@ export function pythonEndpointType(
   }
   return EndpointType.Function;
 }
+
+export function hasPythonSelfArgument(endpointType: EndpointType): boolean {
+  // Instance methods and class methods both use `Argument[self]` for the first parameter. The first
+  // parameter after self is called `Argument[0]`.
+  return (
+    endpointType === EndpointType.Method ||
+    endpointType === EndpointType.ClassMethod
+  );
+}