Skip to content

Commit f4eed4d

Browse files
authored
Merge pull request #3352 from github/koesie10/python-mad-format
Add support for Python in the model editor
2 parents c9a7c11 + 070af9e commit f4eed4d

8 files changed

Lines changed: 459 additions & 4 deletions

File tree

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,14 @@ import type {
33
ModelsAsDataLanguage,
44
ModelsAsDataLanguagePredicates,
55
} from "./models-as-data";
6+
import { python } from "./python";
67
import { ruby } from "./ruby";
78
import { staticLanguage } from "./static";
89

910
const languages: Partial<Record<QueryLanguage, ModelsAsDataLanguage>> = {
1011
[QueryLanguage.CSharp]: staticLanguage,
1112
[QueryLanguage.Java]: staticLanguage,
13+
[QueryLanguage.Python]: python,
1214
[QueryLanguage.Ruby]: ruby,
1315
};
1416

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
import { parseAccessPathTokens } from "../../shared/access-paths";
2+
import type { MethodDefinition } from "../../method";
3+
import { EndpointType } from "../../method";
4+
5+
const memberTokenRegex = /^Member\[(.+)]$/;
6+
7+
export function parsePythonAccessPath(path: string): {
8+
typeName: string;
9+
methodName: string;
10+
endpointType: EndpointType;
11+
path: string;
12+
} {
13+
const tokens = parseAccessPathTokens(path);
14+
15+
if (tokens.length === 0) {
16+
return {
17+
typeName: "",
18+
methodName: "",
19+
endpointType: EndpointType.Method,
20+
path: "",
21+
};
22+
}
23+
24+
const typeParts = [];
25+
let endpointType = EndpointType.Function;
26+
27+
let remainingTokens: typeof tokens = [];
28+
29+
for (let i = 0; i < tokens.length; i++) {
30+
const token = tokens[i];
31+
const memberMatch = token.text.match(memberTokenRegex);
32+
if (memberMatch) {
33+
typeParts.push(memberMatch[1]);
34+
} else if (token.text === "Instance") {
35+
endpointType = EndpointType.Method;
36+
} else {
37+
remainingTokens = tokens.slice(i);
38+
break;
39+
}
40+
}
41+
42+
const methodName = typeParts.pop() ?? "";
43+
const typeName = typeParts.join(".");
44+
const remainingPath = remainingTokens.map((token) => token.text).join(".");
45+
46+
return {
47+
methodName,
48+
typeName,
49+
endpointType,
50+
path: remainingPath,
51+
};
52+
}
53+
54+
export function pythonMethodSignature(typeName: string, methodName: string) {
55+
return `${typeName}#${methodName}`;
56+
}
57+
58+
function pythonTypePath(typeName: string) {
59+
if (typeName === "") {
60+
return "";
61+
}
62+
63+
return typeName
64+
.split(".")
65+
.map((part) => `Member[${part}]`)
66+
.join(".");
67+
}
68+
69+
export function pythonMethodPath(
70+
typeName: string,
71+
methodName: string,
72+
endpointType: EndpointType,
73+
) {
74+
if (methodName === "") {
75+
return pythonTypePath(typeName);
76+
}
77+
78+
const typePath = pythonTypePath(typeName);
79+
80+
let result = typePath;
81+
if (typePath !== "" && endpointType === EndpointType.Method) {
82+
result += ".Instance";
83+
}
84+
85+
if (result !== "") {
86+
result += ".";
87+
}
88+
89+
result += `Member[${methodName}]`;
90+
91+
return result;
92+
}
93+
94+
export function pythonPath(
95+
typeName: string,
96+
methodName: string,
97+
endpointType: EndpointType,
98+
path: string,
99+
) {
100+
const methodPath = pythonMethodPath(typeName, methodName, endpointType);
101+
if (methodPath === "") {
102+
return path;
103+
}
104+
105+
if (path === "") {
106+
return methodPath;
107+
}
108+
109+
return `${methodPath}.${path}`;
110+
}
111+
112+
export function pythonEndpointType(
113+
method: Omit<MethodDefinition, "endpointType">,
114+
): EndpointType {
115+
if (method.methodParameters.startsWith("(self,")) {
116+
return EndpointType.Method;
117+
}
118+
return EndpointType.Function;
119+
}
Lines changed: 207 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,207 @@
1+
import type { ModelsAsDataLanguage } from "../models-as-data";
2+
import { sharedExtensiblePredicates, sharedKinds } from "../shared";
3+
import { Mode } from "../../shared/mode";
4+
import type { MethodArgument } from "../../method";
5+
import { EndpointType, getArgumentsList } from "../../method";
6+
import {
7+
parsePythonAccessPath,
8+
pythonEndpointType,
9+
pythonMethodPath,
10+
pythonMethodSignature,
11+
pythonPath,
12+
} from "./access-paths";
13+
14+
export const python: ModelsAsDataLanguage = {
15+
availableModes: [Mode.Framework],
16+
createMethodSignature: ({ typeName, methodName }) =>
17+
`${typeName}#${methodName}`,
18+
endpointTypeForEndpoint: (method) => pythonEndpointType(method),
19+
predicates: {
20+
source: {
21+
extensiblePredicate: sharedExtensiblePredicates.source,
22+
supportedKinds: sharedKinds.source,
23+
supportedEndpointTypes: [EndpointType.Method, EndpointType.Function],
24+
// extensible predicate sourceModel(
25+
// string type, string path, string kind
26+
// );
27+
generateMethodDefinition: (method) => [
28+
method.packageName,
29+
pythonPath(
30+
method.typeName,
31+
method.methodName,
32+
method.endpointType,
33+
method.output,
34+
),
35+
method.kind,
36+
],
37+
readModeledMethod: (row) => {
38+
const packageName = row[0] as string;
39+
const {
40+
typeName,
41+
methodName,
42+
endpointType,
43+
path: output,
44+
} = parsePythonAccessPath(row[1] as string);
45+
return {
46+
type: "source",
47+
output,
48+
kind: row[2] as string,
49+
provenance: "manual",
50+
signature: pythonMethodSignature(typeName, methodName),
51+
endpointType,
52+
packageName,
53+
typeName,
54+
methodName,
55+
methodParameters: "",
56+
};
57+
},
58+
},
59+
sink: {
60+
extensiblePredicate: sharedExtensiblePredicates.sink,
61+
supportedKinds: sharedKinds.sink,
62+
supportedEndpointTypes: [EndpointType.Method, EndpointType.Function],
63+
// extensible predicate sinkModel(
64+
// string type, string path, string kind
65+
// );
66+
generateMethodDefinition: (method) => {
67+
return [
68+
method.packageName,
69+
pythonPath(
70+
method.typeName,
71+
method.methodName,
72+
method.endpointType,
73+
method.input,
74+
),
75+
method.kind,
76+
];
77+
},
78+
readModeledMethod: (row) => {
79+
const packageName = row[0] as string;
80+
const {
81+
typeName,
82+
methodName,
83+
endpointType,
84+
path: input,
85+
} = parsePythonAccessPath(row[1] as string);
86+
return {
87+
type: "sink",
88+
input,
89+
kind: row[2] as string,
90+
provenance: "manual",
91+
signature: pythonMethodSignature(typeName, methodName),
92+
endpointType,
93+
packageName,
94+
typeName,
95+
methodName,
96+
methodParameters: "",
97+
};
98+
},
99+
},
100+
summary: {
101+
extensiblePredicate: sharedExtensiblePredicates.summary,
102+
supportedKinds: sharedKinds.summary,
103+
supportedEndpointTypes: [EndpointType.Method, EndpointType.Function],
104+
// extensible predicate summaryModel(
105+
// string type, string path, string input, string output, string kind
106+
// );
107+
generateMethodDefinition: (method) => [
108+
method.packageName,
109+
pythonMethodPath(
110+
method.typeName,
111+
method.methodName,
112+
method.endpointType,
113+
),
114+
method.input,
115+
method.output,
116+
method.kind,
117+
],
118+
readModeledMethod: (row) => {
119+
const packageName = row[0] as string;
120+
const { typeName, methodName, endpointType, path } =
121+
parsePythonAccessPath(row[1] as string);
122+
if (path !== "") {
123+
throw new Error("Summary path must be a method");
124+
}
125+
return {
126+
type: "summary",
127+
input: row[2] as string,
128+
output: row[3] as string,
129+
kind: row[4] as string,
130+
provenance: "manual",
131+
signature: pythonMethodSignature(typeName, methodName),
132+
endpointType,
133+
packageName,
134+
typeName,
135+
methodName,
136+
methodParameters: "",
137+
};
138+
},
139+
},
140+
neutral: {
141+
extensiblePredicate: sharedExtensiblePredicates.neutral,
142+
supportedKinds: sharedKinds.neutral,
143+
// extensible predicate neutralModel(
144+
// string type, string path, string kind
145+
// );
146+
generateMethodDefinition: (method) => [
147+
method.packageName,
148+
pythonMethodPath(
149+
method.typeName,
150+
method.methodName,
151+
method.endpointType,
152+
),
153+
method.kind,
154+
],
155+
readModeledMethod: (row) => {
156+
const packageName = row[0] as string;
157+
const { typeName, methodName, endpointType, path } =
158+
parsePythonAccessPath(row[1] as string);
159+
if (path !== "") {
160+
throw new Error("Neutral path must be a method");
161+
}
162+
return {
163+
type: "neutral",
164+
kind: row[2] as string,
165+
provenance: "manual",
166+
signature: pythonMethodSignature(typeName, methodName),
167+
endpointType,
168+
packageName,
169+
typeName,
170+
methodName,
171+
methodParameters: "",
172+
};
173+
},
174+
},
175+
},
176+
getArgumentOptions: (method) => {
177+
// Argument and Parameter are equivalent in Python, but we'll use Argument in the model editor
178+
const argumentsList = getArgumentsList(method.methodParameters).map(
179+
(argument, index): MethodArgument => {
180+
if (argument.endsWith(":")) {
181+
return {
182+
path: `Argument[${argument}]`,
183+
label: `Argument[${argument}]`,
184+
};
185+
}
186+
187+
return {
188+
path: `Argument[${index}]`,
189+
label: `Argument[${index}]: ${argument}`,
190+
};
191+
},
192+
);
193+
194+
return {
195+
options: [
196+
{
197+
path: "Argument[self]",
198+
label: "Argument[self]",
199+
},
200+
...argumentsList,
201+
],
202+
// If there are no arguments, we will default to "Argument[self]"
203+
defaultArgumentPath:
204+
argumentsList.length > 0 ? argumentsList[0].path : "Argument[self]",
205+
};
206+
},
207+
};

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ export enum EndpointType {
2828
Class = "class",
2929
Method = "method",
3030
Constructor = "constructor",
31+
Function = "function",
3132
}
3233

3334
export interface MethodDefinition {

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,7 @@ export class ModelEditorModule extends DisposableObject {
214214
queryDir,
215215
language,
216216
this.modelConfig,
217+
initialMode,
217218
);
218219
if (!success) {
219220
await cleanupQueryDir();

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import {
99
} from "./model-editor-queries";
1010
import type { CodeQLCliServer } from "../codeql-cli/cli";
1111
import type { ModelConfig } from "../config";
12-
import { Mode } from "./shared/mode";
12+
import type { Mode } from "./shared/mode";
1313
import type { NotificationLogger } from "../common/logging";
1414

1515
/**
@@ -31,6 +31,7 @@ import type { NotificationLogger } from "../common/logging";
3131
* @param queryDir The directory to set up.
3232
* @param language The language to use for the queries.
3333
* @param modelConfig The model config to use.
34+
* @param initialMode The initial mode to use to check the existence of the queries.
3435
* @returns true if the setup was successful, false otherwise.
3536
*/
3637
export async function setUpPack(
@@ -39,6 +40,7 @@ export async function setUpPack(
3940
queryDir: string,
4041
language: QueryLanguage,
4142
modelConfig: ModelConfig,
43+
initialMode: Mode,
4244
): Promise<boolean> {
4345
// Download the required query packs
4446
await cliServer.packDownload([`codeql/${language}-queries`]);
@@ -48,7 +50,7 @@ export async function setUpPack(
4850
const applicationModeQuery = await resolveEndpointsQuery(
4951
cliServer,
5052
language,
51-
Mode.Application,
53+
initialMode,
5254
[],
5355
[],
5456
);

0 commit comments

Comments
 (0)