Handle multiple imports.

Update tests.

Author: khatchad

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

@@ -1768,7 +1768,12 @@ public void testAdd105()
   @Test
   public void testAdd106()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
-    test("tf2_test_add106.py", "add", 2, 2, Map.of(2, Set.of(MNIST_INPUT), 3, Set.of(MNIST_INPUT)));
+    test(
+        "tf2_test_add106.py",
+        "add",
+        2,
+        2,
+        Map.of(2, Set.of(TENSOR_4_FLOAT32), 3, Set.of(TENSOR_4_FLOAT32)));
   }
 
   @Test

com.ibm.wala.cast.python.ml/source/com/ibm/wala/cast/python/ml/client/TensorGenerator.java

@@ -11,6 +11,7 @@
 import static com.ibm.wala.ipa.callgraph.propagation.cfa.CallStringContextSelector.CALL_STRING;
 import static java.util.Arrays.asList;
 import static java.util.Collections.emptyList;
+import static java.util.stream.Collectors.toSet;
 
 import com.ibm.wala.cast.ipa.callgraph.AstPointerKeyFactory;
 import com.ibm.wala.cast.python.ml.types.TensorFlowTypes;
@@ -43,7 +44,6 @@
 import java.util.EnumSet;
 import java.util.List;
 import java.util.Map.Entry;
-import java.util.Optional;
 import java.util.Set;
 import java.util.logging.Logger;
 
@@ -353,11 +353,11 @@ protected EnumSet<DType> getDTypesFromDTypeArgument(
 
       if (typeReference.equals(TensorFlowTypes.D_TYPE)) {
         // we have a dtype.
-        // let's see if it's float32.
+        // let's see if it's a dtype.
         Set<CGNode> importNodes = builder.getCallGraph().getNodes(IMPORT);
 
-        // find the import node from this file.
-        Optional<CGNode> importNode =
+        // find the import nodes from this file.
+        Set<CGNode> importNodesOfInterest =
             importNodes.stream()
                 .filter(
                     in -> {
@@ -384,38 +384,44 @@ protected EnumSet<DType> getDTypesFromDTypeArgument(
 
                       return method.equals(nodeCS.getMethods()[0]);
                     })
-                .findFirst();
+                .collect(toSet());
 
-        InstanceKey tensorFlowIK =
-            pointerAnalysis
-                .getHeapModel()
-                .getInstanceKeyForAllocation(
-                    importNode.get(), NewSiteReference.make(0, TENSORFLOW));
+        if (importNodesOfInterest.isEmpty())
+          throw new IllegalStateException("No import nodes found for source: " + source + ".");
 
-        // Check dtype literals.
         boolean found = false;
 
-        for (Entry<FieldReference, DType> entry : FIELD_REFERENCE_TO_DTYPE.entrySet()) {
-          FieldReference fieldRef = entry.getKey();
-          DType dtype = entry.getValue();
-          IField field = builder.getClassHierarchy().resolveField(fieldRef);
-
-          PointerKey pk =
-              pointerAnalysis.getHeapModel().getPointerKeyForInstanceField(tensorFlowIK, field);
-
-          for (InstanceKey ik : pointerAnalysis.getPointsToSet(pk))
-            if (ik.equals(instanceKey)) {
-              ret.add(dtype);
-              LOGGER.info(
-                  "Found dtype: "
-                      + dtype
-                      + " for source: "
-                      + source
-                      + " from dType: "
-                      + instanceKey
-                      + ".");
-              found = true;
-            }
+        for (CGNode importNode : importNodesOfInterest) {
+          LOGGER.fine("Found import node of interest: " + importNode + ".");
+
+          InstanceKey tensorFlowIK =
+              pointerAnalysis
+                  .getHeapModel()
+                  .getInstanceKeyForAllocation(importNode, NewSiteReference.make(0, TENSORFLOW));
+
+          // Check dtype literals.
+          for (Entry<FieldReference, DType> entry : FIELD_REFERENCE_TO_DTYPE.entrySet()) {
+            FieldReference fieldRef = entry.getKey();
+            DType dtype = entry.getValue();
+            IField field = builder.getClassHierarchy().resolveField(fieldRef);
+
+            PointerKey pk =
+                pointerAnalysis.getHeapModel().getPointerKeyForInstanceField(tensorFlowIK, field);
+
+            for (InstanceKey ik : pointerAnalysis.getPointsToSet(pk))
+              if (ik.equals(instanceKey)) {
+                ret.add(dtype);
+                LOGGER.info(
+                    "Found dtype: "
+                        + dtype
+                        + " for source: "
+                        + source
+                        + " from dType: "
+                        + instanceKey
+                        + ".");
+                found = true;
+              }
+          }
         }
 
         if (!found) throw new IllegalStateException("Unknown dtype: " + instanceKey + ".");