feat(deducer): support searching for the constructor call node corresponding to the caller of a special method call

components/deducers/python-pyright/src/index.ts

@@ -1,73 +1,129 @@
 import * as path from "path";
+import * as fs from "fs-extra";
 import { Uri } from "pyright-internal/dist/common/uri/uri";
-import { DeduceResult, Deducer } from "@plutolang/base/core";
+import { LogLevel } from "pyright-internal/dist/common/console";
 import { Program } from "pyright-internal/dist/analyzer/program";
+import { TypeCategory } from "pyright-internal/dist/analyzer/types";
 import { SourceFile } from "pyright-internal/dist/analyzer/sourceFile";
-import { LogLevel } from "pyright-internal/dist/common/console";
-
+import { TypeEvaluator } from "pyright-internal/dist/analyzer/typeEvaluatorTypes";
+import { ArgumentCategory, CallNode } from "pyright-internal/dist/parser/parseNodes";
+import { DeduceResult, Deducer } from "@plutolang/base/core";
 import * as TextUtils from "./text-utils";
-import * as ProgramUtils from "./program-utils";
 import * as TypeConsts from "./type-consts";
+import * as ProgramUtils from "./program-utils";
+import * as ScopeUtils from "./scope-utils";
 import { TypeSearcher } from "./type-searcher";
-import { ArgumentCategory, CallNode } from "pyright-internal/dist/parser/parseNodes";
-import { TypeEvaluator } from "pyright-internal/dist/analyzer/typeEvaluatorTypes";
-import { TypeCategory } from "pyright-internal/dist/analyzer/types";
+import { SpecialNodeMap } from "./special-node-map";
 import { Value, ValueEvaluator } from "./value-evaluator";
+import { ResourceObjectTracker } from "./resource-object-tracker";
 
 export default class PyrightDeducer extends Deducer {
   //eslint-disable-next-line @typescript-eslint/no-var-requires
   public readonly name = require(path.join(__dirname, "../package.json")).name;
   //eslint-disable-next-line @typescript-eslint/no-var-requires
   public readonly version = require(path.join(__dirname, "../package.json")).version;
 
+  private sepcialNodeMap?: SpecialNodeMap<CallNode>;
+
   public deduce(entrypoints: string[]): Promise<DeduceResult> {
+    if (entrypoints.length === 0) {
+      throw new Error("No entrypoints provided.");
+    }
+    if (entrypoints.length > 1) {
+      throw new Error("Only one entrypoint is supported, currently.");
+    }
+    // Check if all the entrypoint files exist.
+    for (const filepath of entrypoints) {
+      if (!fs.existsSync(filepath)) {
+        throw new Error(`File not found: ${filepath}`);
+      }
+    }
+
     const program = ProgramUtils.createProgram({
       logLevel: LogLevel.Warn,
       extraPaths: [
         path.resolve(__dirname, "../../../../packages/base-py"),
         path.resolve(__dirname, "../../../../packages/pluto-py"),
       ],
     });
-
     const fileUris = entrypoints.map((name) => Uri.file(name));
     program.setTrackedFiles(fileUris);
-
+    // Wait for the analysis to complete
     // eslint-disable-next-line no-empty
     while (program.analyze()) {}
 
     const sourceFile = program.getSourceFile(fileUris[0])!;
-    doTypeSearch(program, sourceFile);
+
+    this.sepcialNodeMap = this.getSecpialNodes(program, sourceFile);
+    const constructNodes = this.sepcialNodeMap.getNodesByType(TypeConsts.IRESOURCE_FULL_NAME);
+    const infraApiNodes = this.sepcialNodeMap.getNodesByType(
+      TypeConsts.IRESOURCE_INFRA_API_FULL_NAME
+    );
+    for (const node of [constructNodes, infraApiNodes].flat()) {
+      if (node && !ScopeUtils.inGlobalScope(node, sourceFile)) {
+        throw new Error(
+          "All constructor and infrastructre API calls related to pluto resource types should be in global scope. We will relax this restriction in the future."
+        );
+      }
+    }
+
+    const tracker = new ResourceObjectTracker(program.evaluator!, this.sepcialNodeMap);
+
+    const specialTypes = this.sepcialNodeMap.getSpicalTypes();
+    console.log(specialTypes.length, "types of special nodes found.");
+    for (const specialType of specialTypes) {
+      const nodes = this.sepcialNodeMap.getNodesByType(specialType)!;
+
+      console.log("Special Node:", specialType);
+      nodes.forEach((node) => {
+        console.log("/--------------------\\");
+        console.log("|", TextUtils.getTextOfNode(node, sourceFile));
+
+        if (specialType === TypeConsts.IRESOURCE_FULL_NAME) {
+          console.log("| NodeID:", node.id);
+        }
+
+        if (
+          specialType === TypeConsts.IRESOURCE_INFRA_API_FULL_NAME ||
+          specialType === TypeConsts.IRESOURCE_CLIENT_API_FULL_NAME
+        ) {
+          const constuctNode = tracker.getConstructNodeForApiCall(node, sourceFile);
+          if (!constuctNode) {
+            const nodeText = TextUtils.getTextOfNode(node, sourceFile);
+            throw new Error(`No related node found for node '${nodeText}'.`);
+          }
+          console.log("| Related Node ID: ", constuctNode.id);
+        }
+
+        if (
+          specialType === TypeConsts.IRESOURCE_FULL_NAME ||
+          specialType === TypeConsts.IRESOURCE_INFRA_API_FULL_NAME
+        ) {
+          getArgumentValue(node, sourceFile, program.evaluator!);
+          // console.log(inGlobalScope(node, sourceFile));
+        }
+        console.log("\\--------------------/\n\n");
+      });
+    }
 
     program.dispose();
     return {} as any;
   }
-}
 
-function doTypeSearch(program: Program, sourceFile: SourceFile) {
-  const parseResult = sourceFile.getParseResults();
-  if (!parseResult) {
-    throw new Error("No parse result");
+  /**
+   * Use the TypeSearcher to get the special nodes in the source file.
+   */
+  private getSecpialNodes(program: Program, sourceFile: SourceFile) {
+    const parseResult = sourceFile.getParseResults();
+    if (!parseResult) {
+      throw new Error(`No parse result found in source file '${sourceFile.getUri().key}'.`);
+    }
+    const parseTree = parseResult.parseTree;
+
+    const walker = new TypeSearcher(program.evaluator!, sourceFile);
+    walker.walk(parseTree);
+    return walker.specialNodeMap;
   }
-  const parseTree = parseResult.parseTree;
-
-  const walker = new TypeSearcher(program.evaluator!, sourceFile);
-  walker.walk(parseTree);
-
-  console.log(walker.specialNodeMap.size, "types of special nodes found.");
-  walker.specialNodeMap.forEach((nodes, key) => {
-    console.log("Special Node:", key);
-    nodes.forEach((node) => {
-      console.log("/--------------------\\");
-      console.log("|", TextUtils.getTextOfNode(node, sourceFile));
-      if (
-        key === TypeConsts.IRESOURCE_FULL_NAME ||
-        key === TypeConsts.IRESOURCE_INFRA_API_FULL_NAME
-      ) {
-        getArgumentValue(node, sourceFile, program.evaluator!);
-      }
-      console.log("\\--------------------/\n\n");
-    });
-  });
 }
 
 function getArgumentValue(

components/deducers/python-pyright/src/resource-object-tracker.ts

@@ -0,0 +1,162 @@
+import assert from "node:assert";
+import { TypeCategory } from "pyright-internal/dist/analyzer/types";
+import { SourceFile } from "pyright-internal/dist/analyzer/sourceFile";
+import { DeclarationType } from "pyright-internal/dist/analyzer/declaration";
+import { TypeEvaluator } from "pyright-internal/dist/analyzer/typeEvaluatorTypes";
+import { CallNode, NameNode, ParseNodeType } from "pyright-internal/dist/parser/parseNodes";
+import * as TextUtils from "./text-utils";
+import * as TypeConsts from "./type-consts";
+import { SpecialNodeMap } from "./special-node-map";
+
+export class ResourceObjectTracker {
+  private readonly cache = new Map<number, CallNode>();
+
+  constructor(
+    private readonly typeEvaluator: TypeEvaluator,
+    private readonly sepcialNodeMap: SpecialNodeMap<CallNode>
+  ) {}
+
+  /**
+   * Get the node that constructs the resource object the API call is on. The target node should be
+   * within the special node map created by the TypeSearcher.
+   * @param callNode - The node of the API calls, including the client API, the infrastructure API
+   * and the captured property.
+   * @param sourceFile - The source file that the API call is in.
+   * @returns - The node of the resource object that the API call is on. If the node is not found,
+   * return undefined.
+   */
+  public getConstructNodeForApiCall(callNode: CallNode, sourceFile: SourceFile) {
+    if (this.cache.has(callNode.id)) {
+      return this.cache.get(callNode.id);
+    }
+
+    const apiExpression = callNode.leftExpression;
+    if (apiExpression.nodeType !== ParseNodeType.MemberAccess) {
+      // This call expression isn't directly calling a method on a resource object, but might be
+      // calling a function variable assigned to a method of a resource object. For example:
+      // ```python
+      // func = router.get
+      // func("/path", handler)
+      // ```
+      const nodeText = TextUtils.getTextOfNode(callNode, sourceFile);
+      throw new Error(
+        `Failed to process this expression '${nodeText}'. We currently only support directly calling methods on the resource object; indirect method calls, such as assigning a method to a variable and then invoking it, are not supported.`
+      );
+    }
+
+    let constructNode: CallNode | undefined;
+    const callerNode = apiExpression.leftExpression;
+    switch (callerNode.nodeType) {
+      case ParseNodeType.Name:
+        // The caller is a variable. We need to find the resource object that this variable is
+        // assigned to, and get the node that constructs this resource object.
+        constructNode = this.getConstructNodeByNameNode(callerNode, sourceFile);
+        break;
+      case ParseNodeType.Call:
+        // The caller is a direct constructor call. We attempt to locate this node in the special
+        // node map. If it's located, it means that the resource object being constructed by the
+        // constructor call is the same as the one the API call is made on. If it's not found, an
+        // error occurs, and this function's caller will address it.
+        constructNode = this.getConstructNodeByCallNode(callerNode);
+        break;
+      default:
+        throw new Error(`The caller node type '${callerNode.nodeType}' is not supported.`);
+    }
+
+    if (constructNode) {
+      this.cache.set(callNode.id, constructNode);
+    }
+    return constructNode;
+  }
+
+  /**
+   * The caller variable refers to a resource object. We're attempting to locate the calling node
+   * that creates this resource object.
+   * @param node - The name node refering to the resource object.
+   * @param sourceFile - The source file that the node is in.
+   * @returns - The node that constructs the resource object. If the node is not found, return
+   * undefined.
+   */
+  private getConstructNodeByNameNode(node: NameNode, sourceFile: SourceFile): CallNode | undefined {
+    if (this.cache.has(node.id)) {
+      return this.cache.get(node.id);
+    }
+
+    const symbolTable = sourceFile.getModuleSymbolTable();
+    if (!symbolTable) {
+      throw new Error(`No symbol table found in source file '${sourceFile.getUri().key}'.`);
+    }
+
+    const symbol = symbolTable.get(node.value);
+    if (!symbol) {
+      throw new Error(`No symbol found for node '${node.value}'.`);
+    }
+
+    if (symbol.getDeclarations().length !== 1) {
+      // If there are multiple declarations, we can't determine which one corresponds to the current
+      // node.
+      throw new Error(
+        `Currently, we only support the resource variable only can be assigned once.`
+      );
+    }
+
+    const declaration = symbol.getDeclarations()[0];
+    if (declaration.type !== DeclarationType.Variable) {
+      throw new Error(`The declaration type '${declaration.type}' is not variable.`);
+    }
+
+    // Since we know the node is in an infrastructure API call, the caller must be a resource
+    // object, so we can trust that the inferred type source of this declaration must not be
+    // undefined.
+    assert(declaration.inferredTypeSource, "No inferred type source");
+    const inferredTypeNode = declaration.inferredTypeSource!;
+
+    // The `declaration.inferredTypeSource` is the source node that can be used to infer the type.
+    // It should be either the function call or another resource object. In the latter case, the
+    // current node is an alias variable of the other resource object.
+    let constructNode: CallNode | undefined;
+    switch (inferredTypeNode.nodeType) {
+      case ParseNodeType.Call:
+        // The inferred type source node is a direct constructor call. Search this node from the
+        // special node map created by the TypeSearcher. If it's found, return it; otherwise, it is
+        // an error, and the caller will handle it.
+        constructNode = this.getConstructNodeByCallNode(inferredTypeNode);
+        break;
+      case ParseNodeType.Name:
+        // The inferred type source node is another resource object.
+        constructNode = this.getConstructNodeByNameNode(inferredTypeNode, sourceFile);
+        break;
+      default:
+        throw new Error(
+          `The inferred type source node type '${inferredTypeNode.nodeType}' is not supported.`
+        );
+    }
+
+    if (constructNode) {
+      this.cache.set(node.id, constructNode);
+    }
+    return constructNode;
+  }
+
+  private getConstructNodeByCallNode(callNode: CallNode): CallNode | undefined {
+    if (this.cache.has(callNode.id)) {
+      return this.cache.get(callNode.id);
+    }
+
+    const callType = this.typeEvaluator.getType(callNode.leftExpression);
+    if (!callType || callType.category !== TypeCategory.Class) {
+      throw new Error(
+        `The type of the call node '${callType?.category}' is not a class.We currently only support the variable assigned from a class constructor.`
+      );
+    }
+    const constructNode = this.sepcialNodeMap.getNodeById(
+      callNode.id,
+      TypeConsts.IRESOURCE_FULL_NAME
+    );
+
+    if (constructNode) {
+      this.cache.set(callNode.id, constructNode);
+    }
+    return constructNode;
+  }
+}

components/deducers/python-pyright/src/scope-utils.ts

@@ -0,0 +1,12 @@
+import assert from "node:assert";
+import { getScopeForNode } from "pyright-internal/dist/analyzer/scopeUtils";
+import { SourceFile } from "pyright-internal/dist/analyzer/sourceFile";
+import { CallNode } from "pyright-internal/dist/parser/parseNodes";
+
+export function inGlobalScope(node: CallNode, sourceFile: SourceFile): boolean {
+  const scope = getScopeForNode(node);
+  const parseTree = sourceFile.getParseResults()?.parseTree;
+  assert(parseTree, "No parse tree found in source file.");
+  const globalScope = getScopeForNode(parseTree!);
+  return scope === globalScope;
+}

components/deducers/python-pyright/src/special-node-map.ts

@@ -0,0 +1,37 @@
+import { ParseNode } from "pyright-internal/dist/parser/parseNodes";
+
+/**
+ * A map that stores special nodes. Key is the full qualified name of the special type. Value is an
+ * array of nodes.
+ */
+export class SpecialNodeMap<T extends ParseNode> {
+  private readonly map: Map<string, T[]> = new Map();
+
+  public addNode(key: string, node: T): void {
+    if (!this.map.has(key)) {
+      this.map.set(key, []);
+    }
+    this.map.get(key)!.push(node);
+  }
+
+  public getNodeById(nodeId: number, key?: string): T | undefined {
+    if (key) {
+      const nodes = this.map.get(key);
+      return nodes?.find((node) => node.id === nodeId);
+    }
+
+    for (const nodes of this.map.values()) {
+      const node = nodes.find((n) => n.id === nodeId);
+      if (node) return node;
+    }
+    return undefined;
+  }
+
+  public getNodesByType(specialTypeName: string): T[] | undefined {
+    return this.map.get(specialTypeName);
+  }
+
+  public getSpicalTypes(): string[] {
+    return Array.from(this.map.keys());
+  }
+}