From d40186b95c0b3f958b3052a4b7f82801fe2e223f Mon Sep 17 00:00:00 2001
From: Raffi Khatchadourian <khatchad@hunter.cuny.edu>
Date: Wed, 1 May 2024 11:37:50 -0400
Subject: [PATCH] Decorator fixes.

---
 .../python/ml/test/TestTensorflow2Model.java  | 11 +++++
 .../data/tensorflow.xml                       | 22 ++++++++++
 .../data/tf2_test_decorated_functions.py      | 41 +++++++++++++++++++
 com.ibm.wala.cast.python/data/pytest.xml      | 23 ++++++++++-
 4 files changed, 95 insertions(+), 2 deletions(-)
 create mode 100644 com.ibm.wala.cast.python.test/data/tf2_test_decorated_functions.py

diff --git a/com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflow2Model.java b/com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflow2Model.java
index 8061d9c49..7478bf03f 100644
--- a/com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflow2Model.java
+++ b/com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflow2Model.java
@@ -3691,6 +3691,17 @@ public void testDecoratedMethod13() throws ClassHierarchyException, CancelExcept
     test("tf2_test_decorated_method13.py", "f", 0, 0);
   }
 
+  @Test
+  public void testDecoratedFunctions()
+      throws ClassHierarchyException, CancelException, IOException {
+    test("tf2_test_decorated_functions.py", "dummy_fun", 1, 1, 2);
+    test("tf2_test_decorated_functions.py", "dummy_test", 1, 1, 2);
+    test("tf2_test_decorated_functions.py", "test_function", 1, 1, 2);
+    test("tf2_test_decorated_functions.py", "test_function2", 1, 1, 2);
+    test("tf2_test_decorated_functions.py", "test_function3", 1, 1, 2);
+    test("tf2_test_decorated_functions.py", "test_function4", 1, 1, 2);
+  }
+
   private void test(
       String filename,
       String functionName,
diff --git a/com.ibm.wala.cast.python.ml/data/tensorflow.xml b/com.ibm.wala.cast.python.ml/data/tensorflow.xml
index f49cf34cc..adf31dc52 100644
--- a/com.ibm.wala.cast.python.ml/data/tensorflow.xml
+++ b/com.ibm.wala.cast.python.ml/data/tensorflow.xml
@@ -252,6 +252,12 @@
         <putfield class="LRoot" field="from_row_starts" fieldType="LRoot" ref="RaggedTensor" value="from_row_starts" />
         <new def="from_value_rowids" class="Ltensorflow/functions/from_value_rowids" />
         <putfield class="LRoot" field="from_value_rowids" fieldType="LRoot" ref="RaggedTensor" value="from_value_rowids" />
+        <new def="autograph" class="Lobject" />
+        <putfield class="LRoot" field="autograph" fieldType="LRoot" ref="x" value="autograph" />
+        <putfield class="LRoot" field="experimental" fieldType="LRoot" ref="autograph" value="experimental" />
+        <!-- https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/autograph/experimental/do_not_convert -->
+        <new def="do_not_convert" class="Ltensorflow/class/do_not_convert" />
+        <putfield class="LRoot" field="do_not_convert" fieldType="LRoot" ref="experimental" value="do_not_convert" />
         <return value="x" />
       </method>
     </class>
@@ -291,6 +297,22 @@
           <return value="params" />
         </method>
       </class>
+      <class name="Do_not_convert" allocatable="true">
+        <method name="do" descriptor="()LRoot;" numArgs="2" paramNames="self test">
+          <putfield class="LRoot" field="params" fieldType="LRoot" ref="test" value="self" />
+          <return value="test" />
+        </method>
+      </class>
+      <class name="do_not_convert" allocatable="true">
+        <!-- https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/autograph/experimental/do_not_convert -->
+        <method name="do" descriptor="()LRoot;" numArgs="4" paramNames="self params values extra">
+          <new def="closure" class="Ltensorflow/class/Do_not_convert" />
+          <putfield class="LRoot" field="test" fieldType="LRoot" ref="closure" value="self" />
+          <putfield class="LRoot" field="params" fieldType="LRoot" ref="closure" value="params" />
+          <putfield class="LRoot" field="values" fieldType="LRoot" ref="closure" value="values" />
+          <return value="closure" />
+        </method>
+      </class>
     </package>
     <package name="tensorflow/objects">
       <class name="feature" allocatable="true" />
diff --git a/com.ibm.wala.cast.python.test/data/tf2_test_decorated_functions.py b/com.ibm.wala.cast.python.test/data/tf2_test_decorated_functions.py
new file mode 100644
index 000000000..57cfd06cf
--- /dev/null
+++ b/com.ibm.wala.cast.python.test/data/tf2_test_decorated_functions.py
@@ -0,0 +1,41 @@
+import tensorflow
+import pytest
+import sys
+
+
+@tensorflow.autograph.experimental.do_not_convert
+def dummy_fun(a):
+    pass
+
+
+@pytest.mark.parametrize("test_input,expected", [("3+5", 8), ("2+4", 6), ("6*9", 42)])
+def dummy_test(x, test_input, expected):
+    pass
+
+
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="requires python3.10 or higher")
+def test_function(x):
+    pass
+
+
+@pytest.mark.skip(reason="requires python3.10 or higher")
+def test_function2(x):
+    pass
+
+
+@pytest.mark.skip("requires python3.10 or higher")
+def test_function3(x):
+    pass
+
+
+@pytest.mark.skip
+def test_function4(x):
+    pass
+
+
+dummy_fun(tensorflow.constant(1))
+dummy_test(tensorflow.constant(1), "1", "1")
+test_function(tensorflow.constant(1))
+test_function2(tensorflow.constant(1))
+test_function3(tensorflow.constant(1))
+test_function4(tensorflow.constant(1))
diff --git a/com.ibm.wala.cast.python/data/pytest.xml b/com.ibm.wala.cast.python/data/pytest.xml
index e75bb2814..60d5e27ad 100644
--- a/com.ibm.wala.cast.python/data/pytest.xml
+++ b/com.ibm.wala.cast.python/data/pytest.xml
@@ -12,6 +12,8 @@
         <putfield class="LRoot" field="parametrize" fieldType="LRoot" ref="mark" value="parametrize" />
         <new def="skip" class="Lpytest/class/skip" />
         <putfield class="LRoot" field="skip" fieldType="LRoot" ref="mark" value="skip" />
+        <new def="skipif" class="Lpytest/class/skipif" />
+        <putfield class="LRoot" field="skipif" fieldType="LRoot" ref="mark" value="skipif" />
         <return value="x" />
       </method>
     </class>
@@ -39,10 +41,27 @@
       </class>
       <class name="skip" allocatable="true">
         <!-- https://docs.pytest.org/en/8.2.x/reference/reference.html#pytest.mark.skip -->
-        <method name="do" descriptor="()LRoot;" numArgs="2" paramNames="self reason">
+        <method name="do" descriptor="()LRoot;" numArgs="4" paramNames="self params values extra">
           <new def="closure" class="Lpytest/class/Skip" />
           <putfield class="LRoot" field="test" fieldType="LRoot" ref="closure" value="self" />
-          <putfield class="LRoot" field="params" fieldType="LRoot" ref="closure" value="reason" />
+          <putfield class="LRoot" field="params" fieldType="LRoot" ref="closure" value="params" />
+          <putfield class="LRoot" field="values" fieldType="LRoot" ref="closure" value="values" />
+          <return value="closure" />
+        </method>
+      </class>
+      <class name="Skipif" allocatable="true">
+        <method name="do" descriptor="()LRoot;" numArgs="2" paramNames="self test">
+          <putfield class="LRoot" field="params" fieldType="LRoot" ref="test" value="self" />
+          <return value="test" />
+        </method>
+      </class>
+      <class name="skipif" allocatable="true">
+        <!-- https://docs.pytest.org/en/8.2.x/reference/reference.html#pytest.mark.skipif -->
+        <method name="do" descriptor="()LRoot;" numArgs="4" paramNames="self params values extra">
+          <new def="closure" class="Lpytest/class/Skipif" />
+          <putfield class="LRoot" field="test" fieldType="LRoot" ref="closure" value="self" />
+          <putfield class="LRoot" field="params" fieldType="LRoot" ref="closure" value="params" />
+          <putfield class="LRoot" field="values" fieldType="LRoot" ref="closure" value="values" />
           <return value="closure" />
         </method>
       </class>