Skip to content

Commit

Permalink
Support static methods (#186)
Browse files Browse the repository at this point in the history
  • Loading branch information
khatchad authored Apr 25, 2024
1 parent 6e73bc6 commit 106d708
Show file tree
Hide file tree
Showing 18 changed files with 512 additions and 25 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,10 @@
*****************************************************************************/
package com.ibm.wala.cast.python.parser;

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;

import com.ibm.wala.cast.ir.translator.AbstractClassEntity;
import com.ibm.wala.cast.ir.translator.AbstractCodeEntity;
import com.ibm.wala.cast.ir.translator.AbstractFieldEntity;
Expand All @@ -19,6 +23,7 @@
import com.ibm.wala.cast.python.loader.DynamicAnnotatableEntity;
import com.ibm.wala.cast.python.types.PythonTypes;
import com.ibm.wala.cast.tree.CAst;
import com.ibm.wala.cast.tree.CAstAnnotation;
import com.ibm.wala.cast.tree.CAstEntity;
import com.ibm.wala.cast.tree.CAstNode;
import com.ibm.wala.cast.tree.CAstQualifier;
Expand Down Expand Up @@ -49,6 +54,7 @@
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.Map;
Expand Down Expand Up @@ -1155,6 +1161,35 @@ public String toString() {
}
;

Collection<CAstAnnotation> annotations = new ArrayList<>();

for (CAstNode node : dynamicAnnotations) {
CAstAnnotation cAstAnnotation =
new CAstAnnotation() {
@Override
public CAstType getType() {
return PythonTypes.CAST_DYNAMIC_ANNOTATION;
}

@Override
public Map<String, Object> getArguments() {
Map<String, Object> map = new HashMap<>();
map.put(DYNAMIC_ANNOTATION_KEY, node);
return map;
}

@Override
public String toString() {
return this.getArguments().getOrDefault(DYNAMIC_ANNOTATION_KEY, this).toString();
}
};

annotations.add(cAstAnnotation);
}

boolean staticMethod =
getNameStream(annotations).anyMatch(s -> s.equals(STATIC_METHOD_ANNOTATION_NAME));

CAstType functionType;
boolean isMethod =
context.entity().getKind() == CAstEntity.TYPE_ENTITY
Expand All @@ -1170,7 +1205,7 @@ public CAstType getDeclaringType() {

@Override
public boolean isStatic() {
return false;
return staticMethod;
}
}
;
Expand Down Expand Up @@ -1199,14 +1234,25 @@ class PythonCodeEntity extends AbstractCodeEntity
implements PythonGlobalsEntity, DynamicAnnotatableEntity {
private final java.util.Set<String> downwardGlobals;

private final Collection<CAstAnnotation> annotations;

@Override
public Iterable<CAstNode> dynamicAnnotations() {
return dynamicAnnotations;
}

protected PythonCodeEntity(CAstType type, java.util.Set<String> downwardGlobals) {
protected PythonCodeEntity(
CAstType type,
java.util.Set<String> downwardGlobals,
Collection<CAstAnnotation> annotations) {
super(type);
this.downwardGlobals = downwardGlobals;
this.annotations = annotations;
}

@Override
public Collection<CAstAnnotation> getAnnotations() {
return this.annotations;
}

@Override
Expand All @@ -1217,8 +1263,10 @@ public int getKind() {
@Override
public CAstNode getAST() {
if (function instanceof FunctionDef) {
if (isMethod) {
// Only add object metadata for non-static methods.
if (isMethod && !staticMethod) {
CAst Ast = PythonParser.this.Ast;

CAstNode[] newNodes = new CAstNode[nodes.length + 2];
System.arraycopy(nodes, 0, newNodes, 2, nodes.length);

Expand Down Expand Up @@ -1306,7 +1354,7 @@ public java.util.Set<String> downwardGlobals() {

java.util.Set<String> downwardGlobals = HashSetFactory.make();

PythonCodeEntity fun = new PythonCodeEntity(functionType, downwardGlobals);
PythonCodeEntity fun = new PythonCodeEntity(functionType, downwardGlobals, annotations);

PythonParser.FunctionContext child =
new PythonParser.FunctionContext(context, fun, downwardGlobals, function);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3314,6 +3314,192 @@ public void testModule53()
expectedTensorParameterValueNumbers);
}

@Test
public void testStaticMethod() throws ClassHierarchyException, CancelException, IOException {
test("tf2_test_static_method.py", "MyClass.the_static_method", 1, 1, 2);
}

@Test
public void testStaticMethod2() throws ClassHierarchyException, CancelException, IOException {
test("tf2_test_static_method2.py", "MyClass.the_static_method", 1, 1, 2);
}

@Test
public void testStaticMethod3() throws ClassHierarchyException, CancelException, IOException {
test("tf2_test_static_method3.py", "MyClass.the_static_method", 1, 1, 2);
}

@Test
public void testStaticMethod4() throws ClassHierarchyException, CancelException, IOException {
test("tf2_test_static_method4.py", "MyClass.the_static_method", 1, 1, 2);
}

@Test
public void testStaticMethod5() throws ClassHierarchyException, CancelException, IOException {
int expectNumberofTensorParameters;
int expectedNumberOfTensorVariables;
int[] expectedTensorParameterValueNumbers;

// Static methods are only supported for Jython3.
if (usesJython3Testing()) {
expectNumberofTensorParameters = 1;
expectedNumberOfTensorVariables = 1;
expectedTensorParameterValueNumbers = new int[] {2};
} 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_static_method5.py",
"MyClass.the_static_method",
expectNumberofTensorParameters,
expectedNumberOfTensorVariables,
expectedTensorParameterValueNumbers);
}

@Test
public void testStaticMethod6() throws ClassHierarchyException, CancelException, IOException {
int expectNumberofTensorParameters;
int expectedNumberOfTensorVariables;
int[] expectedTensorParameterValueNumbers;

// Static methods are only supported for Jython3.
if (usesJython3Testing()) {
expectNumberofTensorParameters = 1;
expectedNumberOfTensorVariables = 1;
expectedTensorParameterValueNumbers = new int[] {2};
} 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_static_method6.py",
"MyClass.the_static_method",
expectNumberofTensorParameters,
expectedNumberOfTensorVariables,
expectedTensorParameterValueNumbers);
}

@Test
public void testStaticMethod7() throws ClassHierarchyException, CancelException, IOException {
test("tf2_test_static_method7.py", "MyClass.the_static_method", 1, 1, 3);
}

@Test
public void testStaticMethod8() throws ClassHierarchyException, CancelException, IOException {
test("tf2_test_static_method8.py", "MyClass.the_static_method", 1, 1, 3);
}

@Test
public void testStaticMethod9() throws ClassHierarchyException, CancelException, IOException {
int expectNumberofTensorParameters;
int expectedNumberOfTensorVariables;
int[] expectedTensorParameterValueNumbers;

// Static methods are only supported for Jython3.
if (usesJython3Testing()) {
expectNumberofTensorParameters = 2;
expectedNumberOfTensorVariables = 2;
expectedTensorParameterValueNumbers = new int[] {2, 3};
} else {
// NOTE: Remove this case once https://github.com/wala/ML/issues/147 is fixed.
expectNumberofTensorParameters = 1;
expectedNumberOfTensorVariables = 1;
expectedTensorParameterValueNumbers = new int[] {3};
}

test(
"tf2_test_static_method9.py",
"MyClass.the_static_method",
expectNumberofTensorParameters,
expectedNumberOfTensorVariables,
expectedTensorParameterValueNumbers);
}

@Test
public void testStaticMethod10() throws ClassHierarchyException, CancelException, IOException {
int expectNumberofTensorParameters;
int expectedNumberOfTensorVariables;
int[] expectedTensorParameterValueNumbers;

// Static methods are only supported for Jython3.
if (usesJython3Testing()) {
expectNumberofTensorParameters = 2;
expectedNumberOfTensorVariables = 2;
expectedTensorParameterValueNumbers = new int[] {2, 3};
} else {
// NOTE: Remove this case once https://github.com/wala/ML/issues/147 is fixed.
expectNumberofTensorParameters = 1;
expectedNumberOfTensorVariables = 1;
expectedTensorParameterValueNumbers = new int[] {3};
}

test(
"tf2_test_static_method10.py",
"MyClass.the_static_method",
expectNumberofTensorParameters,
expectedNumberOfTensorVariables,
expectedTensorParameterValueNumbers);
}

@Test
public void testStaticMethod11() throws ClassHierarchyException, CancelException, IOException {
int expectNumberofTensorParameters;
int expectedNumberOfTensorVariables;
int[] expectedTensorParameterValueNumbers;

// Static methods are only supported for Jython3.
if (usesJython3Testing()) {
expectNumberofTensorParameters = 1;
expectedNumberOfTensorVariables = 1;
expectedTensorParameterValueNumbers = new int[] {2};
} 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_static_method11.py",
"f",
expectNumberofTensorParameters,
expectedNumberOfTensorVariables,
expectedTensorParameterValueNumbers);
}

@Test
public void testStaticMethod12() throws ClassHierarchyException, CancelException, IOException {
int expectNumberofTensorParameters;
int expectedNumberOfTensorVariables;
int[] expectedTensorParameterValueNumbers;

// Static methods are only supported for Jython3.
if (usesJython3Testing()) {
expectNumberofTensorParameters = 1;
expectedNumberOfTensorVariables = 1;
expectedTensorParameterValueNumbers = new int[] {2};
} 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_static_method12.py",
"f",
expectNumberofTensorParameters,
expectedNumberOfTensorVariables,
expectedTensorParameterValueNumbers);
}

private void test(
String filename,
String functionName,
Expand Down
11 changes: 11 additions & 0 deletions com.ibm.wala.cast.python.test/data/tf2_test_static_method.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
import tensorflow as tf


class MyClass:

@staticmethod
def the_static_method(x):
assert isinstance(x, tf.Tensor)


MyClass.the_static_method(tf.constant(1))
12 changes: 12 additions & 0 deletions com.ibm.wala.cast.python.test/data/tf2_test_static_method10.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
import tensorflow as tf


class MyClass:

@staticmethod
def the_static_method(x, y):
assert isinstance(x, tf.Tensor)
assert isinstance(y, tf.Tensor)


MyClass().the_static_method(tf.constant(1), tf.constant(2))
17 changes: 17 additions & 0 deletions com.ibm.wala.cast.python.test/data/tf2_test_static_method11.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
import tensorflow as tf


def f(x):
assert isinstance(x, tf.Tensor)


class MyClass:

@staticmethod
@tf.function
def the_static_method(x):
assert isinstance(x, tf.Tensor)
f(x)


MyClass().the_static_method(tf.constant(1))
16 changes: 16 additions & 0 deletions com.ibm.wala.cast.python.test/data/tf2_test_static_method12.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
import tensorflow as tf


def f(x):
assert isinstance(x, tf.Tensor)


class MyClass:

@staticmethod
def the_static_method(x):
assert isinstance(x, tf.Tensor)
f(x)


MyClass().the_static_method(tf.constant(1))
12 changes: 12 additions & 0 deletions com.ibm.wala.cast.python.test/data/tf2_test_static_method2.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
import tensorflow as tf


class MyClass:

@staticmethod
@tf.function
def the_static_method(x):
assert isinstance(x, tf.Tensor)


MyClass.the_static_method(tf.constant(1))
11 changes: 11 additions & 0 deletions com.ibm.wala.cast.python.test/data/tf2_test_static_method3.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
import tensorflow as tf


class MyClass:

@tf.function
def the_static_method(x):
assert isinstance(x, tf.Tensor)


MyClass.the_static_method(tf.constant(1))
Loading

0 comments on commit 106d708

Please sign in to comment.