Support class methods

com.ibm.wala.cast.python.jython3/source/com/ibm/wala/cast/python/parser/PythonParser.java

@@ -10,6 +10,7 @@
  *****************************************************************************/
 package com.ibm.wala.cast.python.parser;
 
+import static com.ibm.wala.cast.python.util.Util.CLASS_METHOD_ANNOTATION_NAME;
 import static com.ibm.wala.cast.python.util.Util.DYNAMIC_ANNOTATION_KEY;
 import static com.ibm.wala.cast.python.util.Util.STATIC_METHOD_ANNOTATION_NAME;
 import static com.ibm.wala.cast.python.util.Util.getNameStream;
@@ -1263,8 +1264,12 @@ public int getKind() {
         @Override
         public CAstNode getAST() {
           if (function instanceof FunctionDef) {
-            // Only add object metadata for non-static methods.
-            if (isMethod && !staticMethod) {
+
+            boolean classMethod =
+                getNameStream(annotations).anyMatch(s -> s.equals(CLASS_METHOD_ANNOTATION_NAME));
+
+            // Only add object metadata for non-static and non-class methods.
+            if (isMethod && !staticMethod && !classMethod) {
               CAst Ast = PythonParser.this.Ast;
 
               CAstNode[] newNodes = new CAstNode[nodes.length + 2];

com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflow2Model.java

@@ -3500,6 +3500,115 @@ public void testStaticMethod12() throws ClassHierarchyException, CancelException
         expectedTensorParameterValueNumbers);
   }
 
+  @Test
+  public void testClassMethod() throws ClassHierarchyException, CancelException, IOException {
+    int expectNumberofTensorParameters;
+    int expectedNumberOfTensorVariables;
+    int[] expectedTensorParameterValueNumbers;
+
+    // Class methods are only supported for Jython3.
+    if (usesJython3Testing()) {
+      expectNumberofTensorParameters = 1;
+      expectedNumberOfTensorVariables = 1;
+      expectedTensorParameterValueNumbers = new int[] {3};
+    } else {
+      // NOTE: Remove this case once https://github.com/wala/ML/issues/147 is fixed.
+      expectNumberofTensorParameters = 1;
+      expectedNumberOfTensorVariables = 1;
+      expectedTensorParameterValueNumbers = new int[] {2};
+    }
+
+    test(
+        "tf2_test_class_method.py",
+        "MyClass.the_class_method",
+        expectNumberofTensorParameters,
+        expectedNumberOfTensorVariables,
+        expectedTensorParameterValueNumbers);
+  }
+
+  @Test
+  public void testClassMethod2() throws ClassHierarchyException, CancelException, IOException {
+    test("tf2_test_class_method2.py", "MyClass.the_class_method", 1, 1, 3);
+  }
+
+  @Test
+  public void testClassMethod3() throws ClassHierarchyException, CancelException, IOException {
+    int expectNumberofTensorParameters;
+    int expectedNumberOfTensorVariables;
+    int[] expectedTensorParameterValueNumbers;
+
+    // Class methods are only supported for Jython3.
+    if (usesJython3Testing()) {
+      expectNumberofTensorParameters = 1;
+      expectedNumberOfTensorVariables = 1;
+      expectedTensorParameterValueNumbers = new int[] {3};
+    } else {
+      // NOTE: Remove this case once https://github.com/wala/ML/issues/147 is fixed.
+      expectNumberofTensorParameters = 0;
+      expectedNumberOfTensorVariables = 0;
+      expectedTensorParameterValueNumbers = new int[] {};
+    }
+
+    test(
+        "tf2_test_class_method3.py",
+        "MyClass.f",
+        expectNumberofTensorParameters,
+        expectedNumberOfTensorVariables,
+        expectedTensorParameterValueNumbers);
+  }
+
+  @Test
+  public void testClassMethod4() throws ClassHierarchyException, CancelException, IOException {
+    int expectNumberofTensorParameters;
+    int expectedNumberOfTensorVariables;
+    int[] expectedTensorParameterValueNumbers;
+
+    // Class methods are only supported for Jython3.
+    if (usesJython3Testing()) {
+      expectNumberofTensorParameters = 1;
+      expectedNumberOfTensorVariables = 1;
+      expectedTensorParameterValueNumbers = new int[] {3};
+    } else {
+      // NOTE: Remove this case once https://github.com/wala/ML/issues/147 is fixed.
+      expectNumberofTensorParameters = 0;
+      expectedNumberOfTensorVariables = 0;
+      expectedTensorParameterValueNumbers = new int[] {};
+    }
+
+    test(
+        "tf2_test_class_method4.py",
+        "MyClass.f",
+        expectNumberofTensorParameters,
+        expectedNumberOfTensorVariables,
+        expectedTensorParameterValueNumbers);
+  }
+
+  @Test
+  public void testClassMethod5() throws ClassHierarchyException, CancelException, IOException {
+    int expectNumberofTensorParameters;
+    int expectedNumberOfTensorVariables;
+    int[] expectedTensorParameterValueNumbers;
+
+    // Class methods are only supported for Jython3.
+    if (usesJython3Testing()) {
+      expectNumberofTensorParameters = 1;
+      expectedNumberOfTensorVariables = 1;
+      expectedTensorParameterValueNumbers = new int[] {3};
+    } else {
+      // NOTE: Remove this case once https://github.com/wala/ML/issues/147 is fixed.
+      expectNumberofTensorParameters = 0;
+      expectedNumberOfTensorVariables = 0;
+      expectedTensorParameterValueNumbers = new int[] {};
+    }
+
+    test(
+        "tf2_test_class_method5.py",
+        "MyClass.f",
+        expectNumberofTensorParameters,
+        expectedNumberOfTensorVariables,
+        expectedTensorParameterValueNumbers);
+  }
+
   private void test(
       String filename,
       String functionName,

com.ibm.wala.cast.python.test/data/tf2_test_class_method.py

@@ -0,0 +1,11 @@
+import tensorflow as tf
+
+
+class MyClass:
+
+  @classmethod
+  def the_class_method(cls, x):
+    assert isinstance(x, tf.Tensor)
+
+
+MyClass.the_class_method(tf.constant(1))

com.ibm.wala.cast.python.test/data/tf2_test_class_method2.py

@@ -0,0 +1,11 @@
+import tensorflow as tf
+
+
+class MyClass:
+
+  @classmethod
+  def the_class_method(cls, x):
+    assert isinstance(x, tf.Tensor)
+
+
+MyClass().the_class_method(tf.constant(1))

com.ibm.wala.cast.python.test/data/tf2_test_class_method3.py

@@ -0,0 +1,15 @@
+import tensorflow as tf
+
+
+class MyClass:
+
+  def f(x):
+    assert isinstance(x, tf.Tensor)
+
+  @classmethod
+  def the_class_method(cls, x):
+    assert isinstance(x, tf.Tensor)
+    cls.f(x)
+
+
+MyClass().the_class_method(tf.constant(1))

com.ibm.wala.cast.python.test/data/tf2_test_class_method4.py

@@ -0,0 +1,15 @@
+import tensorflow as tf
+
+
+class MyClass:
+
+  def f(x):
+    assert isinstance(x, tf.Tensor)
+
+  @classmethod
+  def the_class_method(cls, x):
+    assert isinstance(x, tf.Tensor)
+    cls.f(x)
+
+
+MyClass.the_class_method(tf.constant(1))

com.ibm.wala.cast.python.test/data/tf2_test_class_method5.py

@@ -0,0 +1,15 @@
+import tensorflow as tf
+
+
+class MyClass:
+
+  def f(x):
+    assert isinstance(x, tf.Tensor)
+
+  @classmethod
+  def the_class_method(cls, x):
+    assert isinstance(x, tf.Tensor)
+    cls.f(x)
+
+
+MyClass.the_class_method(tf.constant(1))

com.ibm.wala.cast.python/source/com/ibm/wala/cast/python/client/PythonAnalysisEngine.java

@@ -6,10 +6,11 @@
 import com.ibm.wala.cast.ipa.callgraph.AstContextInsensitiveSSAContextInterpreter;
 import com.ibm.wala.cast.ir.ssa.AstIRFactory;
 import com.ibm.wala.cast.loader.AstDynamicField;
+import com.ibm.wala.cast.python.ipa.callgraph.PythonClassMethodTrampolineTargetSelector;
 import com.ibm.wala.cast.python.ipa.callgraph.PythonConstructorTargetSelector;
+import com.ibm.wala.cast.python.ipa.callgraph.PythonInstanceMethodTrampolineTargetSelector;
 import com.ibm.wala.cast.python.ipa.callgraph.PythonSSAPropagationCallGraphBuilder;
 import com.ibm.wala.cast.python.ipa.callgraph.PythonScopeMappingInstanceKeys;
-import com.ibm.wala.cast.python.ipa.callgraph.PythonTrampolineTargetSelector;
 import com.ibm.wala.cast.python.ipa.summaries.BuiltinFunctions;
 import com.ibm.wala.cast.python.ipa.summaries.PythonComprehensionTrampolines;
 import com.ibm.wala.cast.python.ipa.summaries.PythonSuper;
@@ -299,9 +300,10 @@ public boolean isReferenceType() {
 
   protected void addBypassLogic(IClassHierarchy cha, AnalysisOptions options) {
     options.setSelector(
-        new PythonTrampolineTargetSelector<T>(
-            new PythonConstructorTargetSelector(
-                new PythonComprehensionTrampolines(options.getMethodTargetSelector())),
+        new PythonInstanceMethodTrampolineTargetSelector<T>(
+            new PythonClassMethodTrampolineTargetSelector<T>(
+                new PythonConstructorTargetSelector(
+                    new PythonComprehensionTrampolines(options.getMethodTargetSelector()))),
             this));
 
     BuiltinFunctions builtins = new BuiltinFunctions(cha);

com.ibm.wala.cast.python/source/com/ibm/wala/cast/python/ipa/callgraph/PythonClassMethodTrampolineTargetSelector.java

@@ -0,0 +1,131 @@
+/******************************************************************************
+ * Copyright (c) 2018 IBM Corporation.
+ * All rights reserved. This program and the accompanying materials
+ * are made available under the terms of the Eclipse Public License v1.0
+ * which accompanies this distribution, and is available at
+ * http://www.eclipse.org/legal/epl-v10.html
+ *
+ * Contributors:
+ *     IBM Corporation - initial API and implementation
+ *****************************************************************************/
+package com.ibm.wala.cast.python.ipa.callgraph;
+
+import static com.ibm.wala.cast.python.types.Util.getGlobalName;
+import static com.ibm.wala.cast.python.types.Util.makeGlobalRef;
+import static com.ibm.wala.cast.python.util.Util.isClassMethod;
+
+import com.ibm.wala.cast.ir.ssa.AstGlobalRead;
+import com.ibm.wala.cast.loader.DynamicCallSiteReference;
+import com.ibm.wala.cast.python.ipa.summaries.PythonSummary;
+import com.ibm.wala.cast.python.ir.PythonLanguage;
+import com.ibm.wala.cast.python.ssa.PythonInvokeInstruction;
+import com.ibm.wala.cast.python.types.PythonTypes;
+import com.ibm.wala.classLoader.CallSiteReference;
+import com.ibm.wala.classLoader.IClass;
+import com.ibm.wala.core.util.strings.Atom;
+import com.ibm.wala.ipa.callgraph.CGNode;
+import com.ibm.wala.ipa.callgraph.MethodTargetSelector;
+import com.ibm.wala.ipa.cha.IClassHierarchy;
+import com.ibm.wala.ssa.SSAInstructionFactory;
+import com.ibm.wala.ssa.SSAReturnInstruction;
+import com.ibm.wala.types.FieldReference;
+import com.ibm.wala.util.collections.HashMapFactory;
+import com.ibm.wala.util.collections.Pair;
+import java.util.Map;
+import java.util.logging.Logger;
+
+/**
+ * A trampoline for <a href="https://docs.python.org/3/library/functions.html#classmethod">class
+ * methods</a> that are not called using object instances.
+ *
+ * @author <a href="mailto:[email protected]">Raffi Khatchadourian</a>
+ */
+public class PythonClassMethodTrampolineTargetSelector<T>
+    extends PythonMethodTrampolineTargetSelector<T> {
+
+  protected static final Logger LOGGER =
+      Logger.getLogger(PythonClassMethodTrampolineTargetSelector.class.getName());
+
+  public PythonClassMethodTrampolineTargetSelector(MethodTargetSelector base) {
+    super(base);
+  }
+
+  @Override
+  protected boolean shouldProcess(CGNode caller, CallSiteReference site, IClass receiver) {
+    IClassHierarchy cha = receiver.getClassHierarchy();
+
+    // Are we calling a class method?
+    boolean classMethodReceiver = isClassMethod(receiver);
+
+    // Is the caller a trampoline?
+    boolean trampoline =
+        caller
+            .getMethod()
+            .getSelector()
+            .getName()
+            .startsWith(Atom.findOrCreateAsciiAtom("trampoline"));
+
+    return classMethodReceiver
+        && !cha.isSubclassOf(receiver, cha.lookupClass(PythonTypes.trampoline))
+        && !trampoline;
+  }
+
+  @SuppressWarnings("unchecked")
+  @Override
+  protected void populate(
+      PythonSummary x, int v, IClass receiver, PythonInvokeInstruction call, Logger logger) {
+    Map<Integer, Atom> names = HashMapFactory.make();
+    SSAInstructionFactory insts = PythonLanguage.Python.instructionFactory();
+
+    // Read the class from the global scope.
+    String globalName = getGlobalName(receiver.getReference());
+    FieldReference globalRef = makeGlobalRef(receiver.getClassLoader(), globalName);
+    int globalReadRes = v++;
+    int pc = 0;
+
+    x.addStatement(new AstGlobalRead(pc++, globalReadRes, globalRef));
+
+    int getInstRes = v++;
+
+    // Read the field from the class corresponding to the called method.
+    FieldReference method =
+        FieldReference.findOrCreate(
+            PythonTypes.Root, Atom.findOrCreateUnicodeAtom("the_class_method"), PythonTypes.Root);
+
+    x.addStatement(insts.GetInstruction(pc++, getInstRes, globalReadRes, method));
+
+    int i = 0;
+    int paramSize = Math.max(2, call.getNumberOfPositionalParameters() + 1);
+    int[] params = new int[paramSize];
+    params[i++] = getInstRes;
+    params[i++] = globalReadRes;
+
+    for (int j = 1; j < call.getNumberOfPositionalParameters(); j++) params[i++] = j + 1;
+
+    int ki = 0, ji = call.getNumberOfPositionalParameters() + 1;
+    Pair<String, Integer>[] keys = new Pair[0];
+
+    if (call.getKeywords() != null) {
+      keys = new Pair[call.getKeywords().size()];
+
+      for (String k : call.getKeywords()) {
+        names.put(ji, Atom.findOrCreateUnicodeAtom(k));
+        keys[ki++] = Pair.make(k, ji++);
+      }
+    }
+
+    CallSiteReference ref = new DynamicCallSiteReference(call.getCallSite().getDeclaredTarget(), 2);
+
+    int except = v++;
+    int invokeResult = v++;
+
+    x.addStatement(new PythonInvokeInstruction(pc++, invokeResult, except, ref, params, keys));
+    x.addStatement(new SSAReturnInstruction(pc++, invokeResult, false));
+    x.setValueNames(names);
+  }
+
+  @Override
+  protected Logger getLogger() {
+    return LOGGER;
+  }
+}