Merge pull request #3352 from github/koesie10/python-mad-format

Add support for Python in the model editor

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

@@ -3,12 +3,14 @@ import type {
   ModelsAsDataLanguage,
   ModelsAsDataLanguagePredicates,
 } from "./models-as-data";
+import { python } from "./python";
 import { ruby } from "./ruby";
 import { staticLanguage } from "./static";
 
 const languages: Partial<Record<QueryLanguage, ModelsAsDataLanguage>> = {
   [QueryLanguage.CSharp]: staticLanguage,
   [QueryLanguage.Java]: staticLanguage,
+  [QueryLanguage.Python]: python,
   [QueryLanguage.Ruby]: ruby,
 };
 

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

@@ -0,0 +1,119 @@
+import { parseAccessPathTokens } from "../../shared/access-paths";
+import type { MethodDefinition } from "../../method";
+import { EndpointType } from "../../method";
+
+const memberTokenRegex = /^Member\[(.+)]$/;
+
+export function parsePythonAccessPath(path: string): {
+  typeName: string;
+  methodName: string;
+  endpointType: EndpointType;
+  path: string;
+} {
+  const tokens = parseAccessPathTokens(path);
+
+  if (tokens.length === 0) {
+    return {
+      typeName: "",
+      methodName: "",
+      endpointType: EndpointType.Method,
+      path: "",
+    };
+  }
+
+  const typeParts = [];
+  let endpointType = EndpointType.Function;
+
+  let remainingTokens: typeof tokens = [];
+
+  for (let i = 0; i < tokens.length; i++) {
+    const token = tokens[i];
+    const memberMatch = token.text.match(memberTokenRegex);
+    if (memberMatch) {
+      typeParts.push(memberMatch[1]);
+    } else if (token.text === "Instance") {
+      endpointType = EndpointType.Method;
+    } else {
+      remainingTokens = tokens.slice(i);
+      break;
+    }
+  }
+
+  const methodName = typeParts.pop() ?? "";
+  const typeName = typeParts.join(".");
+  const remainingPath = remainingTokens.map((token) => token.text).join(".");
+
+  return {
+    methodName,
+    typeName,
+    endpointType,
+    path: remainingPath,
+  };
+}
+
+export function pythonMethodSignature(typeName: string, methodName: string) {
+  return `${typeName}#${methodName}`;
+}
+
+function pythonTypePath(typeName: string) {
+  if (typeName === "") {
+    return "";
+  }
+
+  return typeName
+    .split(".")
+    .map((part) => `Member[${part}]`)
+    .join(".");
+}
+
+export function pythonMethodPath(
+  typeName: string,
+  methodName: string,
+  endpointType: EndpointType,
+) {
+  if (methodName === "") {
+    return pythonTypePath(typeName);
+  }
+
+  const typePath = pythonTypePath(typeName);
+
+  let result = typePath;
+  if (typePath !== "" && endpointType === EndpointType.Method) {
+    result += ".Instance";
+  }
+
+  if (result !== "") {
+    result += ".";
+  }
+
+  result += `Member[${methodName}]`;
+
+  return result;
+}
+
+export function pythonPath(
+  typeName: string,
+  methodName: string,
+  endpointType: EndpointType,
+  path: string,
+) {
+  const methodPath = pythonMethodPath(typeName, methodName, endpointType);
+  if (methodPath === "") {
+    return path;
+  }
+
+  if (path === "") {
+    return methodPath;
+  }
+
+  return `${methodPath}.${path}`;
+}
+
+export function pythonEndpointType(
+  method: Omit<MethodDefinition, "endpointType">,
+): EndpointType {
+  if (method.methodParameters.startsWith("(self,")) {
+    return EndpointType.Method;
+  }
+  return EndpointType.Function;
+}

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

@@ -0,0 +1,207 @@
+import type { ModelsAsDataLanguage } from "../models-as-data";
+import { sharedExtensiblePredicates, sharedKinds } from "../shared";
+import { Mode } from "../../shared/mode";
+import type { MethodArgument } from "../../method";
+import { EndpointType, getArgumentsList } from "../../method";
+import {
+  parsePythonAccessPath,
+  pythonEndpointType,
+  pythonMethodPath,
+  pythonMethodSignature,
+  pythonPath,
+} from "./access-paths";
+
+export const python: ModelsAsDataLanguage = {
+  availableModes: [Mode.Framework],
+  createMethodSignature: ({ typeName, methodName }) =>
+    `${typeName}#${methodName}`,
+  endpointTypeForEndpoint: (method) => pythonEndpointType(method),
+  predicates: {
+    source: {
+      extensiblePredicate: sharedExtensiblePredicates.source,
+      supportedKinds: sharedKinds.source,
+      supportedEndpointTypes: [EndpointType.Method, EndpointType.Function],
+      // extensible predicate sourceModel(
+      //   string type, string path, string kind
+      // );
+      generateMethodDefinition: (method) => [
+        method.packageName,
+        pythonPath(
+          method.typeName,
+          method.methodName,
+          method.endpointType,
+          method.output,
+        ),
+        method.kind,
+      ],
+      readModeledMethod: (row) => {
+        const packageName = row[0] as string;
+        const {
+          typeName,
+          methodName,
+          endpointType,
+          path: output,
+        } = parsePythonAccessPath(row[1] as string);
+        return {
+          type: "source",
+          output,
+          kind: row[2] as string,
+          provenance: "manual",
+          signature: pythonMethodSignature(typeName, methodName),
+          endpointType,
+          packageName,
+          typeName,
+          methodName,
+          methodParameters: "",
+        };
+      },
+    },
+    sink: {
+      extensiblePredicate: sharedExtensiblePredicates.sink,
+      supportedKinds: sharedKinds.sink,
+      supportedEndpointTypes: [EndpointType.Method, EndpointType.Function],
+      // extensible predicate sinkModel(
+      //   string type, string path, string kind
+      // );
+      generateMethodDefinition: (method) => {
+        return [
+          method.packageName,
+          pythonPath(
+            method.typeName,
+            method.methodName,
+            method.endpointType,
+            method.input,
+          ),
+          method.kind,
+        ];
+      },
+      readModeledMethod: (row) => {
+        const packageName = row[0] as string;
+        const {
+          typeName,
+          methodName,
+          endpointType,
+          path: input,
+        } = parsePythonAccessPath(row[1] as string);
+        return {
+          type: "sink",
+          input,
+          kind: row[2] as string,
+          provenance: "manual",
+          signature: pythonMethodSignature(typeName, methodName),
+          endpointType,
+          packageName,
+          typeName,
+          methodName,
+          methodParameters: "",
+        };
+      },
+    },
+    summary: {
+      extensiblePredicate: sharedExtensiblePredicates.summary,
+      supportedKinds: sharedKinds.summary,
+      supportedEndpointTypes: [EndpointType.Method, EndpointType.Function],
+      // extensible predicate summaryModel(
+      //   string type, string path, string input, string output, string kind
+      // );
+      generateMethodDefinition: (method) => [
+        method.packageName,
+        pythonMethodPath(
+          method.typeName,
+          method.methodName,
+          method.endpointType,
+        ),
+        method.input,
+        method.output,
+        method.kind,
+      ],
+      readModeledMethod: (row) => {
+        const packageName = row[0] as string;
+        const { typeName, methodName, endpointType, path } =
+          parsePythonAccessPath(row[1] as string);
+        if (path !== "") {
+          throw new Error("Summary path must be a method");
+        }
+        return {
+          type: "summary",
+          input: row[2] as string,
+          output: row[3] as string,
+          kind: row[4] as string,
+          provenance: "manual",
+          signature: pythonMethodSignature(typeName, methodName),
+          endpointType,
+          packageName,
+          typeName,
+          methodName,
+          methodParameters: "",
+        };
+      },
+    },
+    neutral: {
+      extensiblePredicate: sharedExtensiblePredicates.neutral,
+      supportedKinds: sharedKinds.neutral,
+      // extensible predicate neutralModel(
+      //   string type, string path, string kind
+      // );
+      generateMethodDefinition: (method) => [
+        method.packageName,
+        pythonMethodPath(
+          method.typeName,
+          method.methodName,
+          method.endpointType,
+        ),
+        method.kind,
+      ],
+      readModeledMethod: (row) => {
+        const packageName = row[0] as string;
+        const { typeName, methodName, endpointType, path } =
+          parsePythonAccessPath(row[1] as string);
+        if (path !== "") {
+          throw new Error("Neutral path must be a method");
+        }
+        return {
+          type: "neutral",
+          kind: row[2] as string,
+          provenance: "manual",
+          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
+    const argumentsList = getArgumentsList(method.methodParameters).map(
+      (argument, index): MethodArgument => {
+        if (argument.endsWith(":")) {
+          return {
+            path: `Argument[${argument}]`,
+            label: `Argument[${argument}]`,
+          };
+        }
+
+        return {
+          path: `Argument[${index}]`,
+          label: `Argument[${index}]: ${argument}`,
+        };
+      },
+    );
+
+    return {
+      options: [
+        {
+          path: "Argument[self]",
+          label: "Argument[self]",
+        },
+        ...argumentsList,
+      ],
+      // If there are no arguments, we will default to "Argument[self]"
+      defaultArgumentPath:
+        argumentsList.length > 0 ? argumentsList[0].path : "Argument[self]",
+    };
+  },
+};

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

@@ -28,6 +28,7 @@ export enum EndpointType {
   Class = "class",
   Method = "method",
   Constructor = "constructor",
+  Function = "function",
 }
 
 export interface MethodDefinition {

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

@@ -214,6 +214,7 @@ export class ModelEditorModule extends DisposableObject {
             queryDir,
             language,
             this.modelConfig,
+            initialMode,
           );
           if (!success) {
             await cleanupQueryDir();

extensions/ql-vscode/src/model-editor/model-editor-queries-setup.ts

@@ -9,7 +9,7 @@ import {
 } from "./model-editor-queries";
 import type { CodeQLCliServer } from "../codeql-cli/cli";
 import type { ModelConfig } from "../config";
-import { Mode } from "./shared/mode";
+import type { Mode } from "./shared/mode";
 import type { NotificationLogger } from "../common/logging";
 
 /**
@@ -31,6 +31,7 @@ import type { NotificationLogger } from "../common/logging";
  * @param queryDir The directory to set up.
  * @param language The language to use for the queries.
  * @param modelConfig The model config to use.
+ * @param initialMode The initial mode to use to check the existence of the queries.
  * @returns true if the setup was successful, false otherwise.
  */
 export async function setUpPack(
@@ -39,6 +40,7 @@ export async function setUpPack(
   queryDir: string,
   language: QueryLanguage,
   modelConfig: ModelConfig,
+  initialMode: Mode,
 ): Promise<boolean> {
   // Download the required query packs
   await cliServer.packDownload([`codeql/${language}-queries`]);
@@ -48,7 +50,7 @@ export async function setUpPack(
   const applicationModeQuery = await resolveEndpointsQuery(
     cliServer,
     language,
-    Mode.Application,
+    initialMode,
     [],
     [],
   );

extensions/ql-vscode/test/unit-tests/model-editor/languages/python/access-paths.test.ts

@@ -0,0 +1,109 @@
+import {
+  parsePythonAccessPath,
+  pythonEndpointType,
+  pythonPath,
+} from "../../../../../src/model-editor/languages/python/access-paths";
+import { EndpointType } from "../../../../../src/model-editor/method";
+
+const testCases: Array<{
+  path: string;
+  method: ReturnType<typeof parsePythonAccessPath>;
+}> = [
+  {
+    path: "Member[CommonTokens].Member[Class].Instance.Member[foo]",
+    method: {
+      typeName: "CommonTokens.Class",
+      methodName: "foo",
+      endpointType: EndpointType.Method,
+      path: "",
+    },
+  },
+  {
+    path: "Member[CommonTokens].Member[Class].Instance.Member[foo].Parameter[self]",
+    method: {
+      typeName: "CommonTokens.Class",
+      methodName: "foo",
+      endpointType: EndpointType.Method,
+      path: "Parameter[self]",
+    },
+  },
+  {
+    path: "Member[getSource].ReturnValue",
+    method: {
+      typeName: "",
+      methodName: "getSource",
+      endpointType: EndpointType.Function,
+      path: "ReturnValue",
+    },
+  },
+  {
+    path: "Member[CommonTokens].Member[makePromise].ReturnValue.Awaited",
+    method: {
+      typeName: "CommonTokens",
+      methodName: "makePromise",
+      endpointType: EndpointType.Function,
+      path: "ReturnValue.Awaited",
+    },
+  },
+  {
+    path: "Member[ArgPos].Member[anyParam].Argument[any]",
+    method: {
+      typeName: "ArgPos",
+      methodName: "anyParam",
+      endpointType: EndpointType.Function,
+      path: "Argument[any]",
+    },
+  },
+  {
+    path: "Member[ArgPos].Instance.Member[self_thing].Argument[self]",
+    method: {
+      typeName: "ArgPos",
+      methodName: "self_thing",
+      endpointType: EndpointType.Method,
+      path: "Argument[self]",
+    },
+  },
+];
+
+describe("parsePythonAccessPath", () => {
+  it.each(testCases)("parses $path", ({ path, method }) => {
+    expect(parsePythonAccessPath(path)).toEqual(method);
+  });
+});
+
+describe("pythonPath", () => {
+  it.each(testCases)("constructs $path", ({ path, method }) => {
+    expect(
+      pythonPath(
+        method.typeName,
+        method.methodName,
+        method.endpointType,
+        method.path,
+      ),
+    ).toEqual(path);
+  });
+});
+
+describe("pythonEndpointType", () => {
+  it("returns method for a method", () => {
+    expect(
+      pythonEndpointType({
+        packageName: "testlib",
+        typeName: "CommonTokens",
+        methodName: "foo",
+        methodParameters: "(self,a)",
+      }),
+    ).toEqual(EndpointType.Method);
+  });
+
+  it("returns function for a function", () => {
+    expect(
+      pythonEndpointType({
+        packageName: "testlib",
+        typeName: "CommonTokens",
+        methodName: "foo",
+        methodParameters: "(a)",
+      }),
+    ).toEqual(EndpointType.Function);
+  });
+});

extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/model-editor-queries.test.ts

@@ -48,7 +48,14 @@ describe("setUpPack", () => {
         llmGeneration: false,
       });
 
-      await setUpPack(cliServer, logger, queryDir, language, modelConfig);
+      await setUpPack(
+        cliServer,
+        logger,
+        queryDir,
+        language,
+        modelConfig,
+        Mode.Application,
+      );
 
       const queryFiles = await readdir(queryDir);
       expect(queryFiles).toEqual(
@@ -106,7 +113,14 @@ describe("setUpPack", () => {
         llmGeneration: false,
       });
 
-      await setUpPack(cliServer, logger, queryDir, language, modelConfig);
+      await setUpPack(
+        cliServer,
+        logger,
+        queryDir,
+        language,
+        modelConfig,
+        Mode.Application,
+      );
 
       const queryFiles = await readdir(queryDir);
       expect(queryFiles.sort()).toEqual(["codeql-pack.yml"].sort());