Adding functions to the python API and testing them (#6477)

This PR adds some functions that are missing in the python API, along with unit tests for them.
Subsequent PR will include additional missing functions.

Also includes a yapf run to reformat the test file.

Co-authored-by: Makai Mann makaim@stanford.edu

2 files changed, 40 insertions, 14 deletions

diff --git a/src/api/python/cvc5.pxd b/src/api/python/cvc5.pxd
index e4c0eb915..b91a9e9c5 100644
--- a/src/api/python/cvc5.pxd
+++ b/src/api/python/cvc5.pxd
@@ -102,6 +102,7 @@ cdef extern from "api/cpp/cvc5.h" namespace "cvc5::api":
         Kind getKind() except +
         Sort getSort() except +
         bint isNull() except +
+        bint isIndexed() except +
         T getIndices[T]() except +
         string toString() except +
 
@@ -182,6 +183,7 @@ cdef extern from "api/cpp/cvc5.h" namespace "cvc5::api":
         Term mkFalse() except +
         Term mkBoolean(bint val) except +
         Term mkPi() except +
+        Term mkInteger(const uint64_t i) except +
         Term mkInteger(const string& s) except +
         Term mkReal(const string& s) except +
         Term mkRegexpEmpty() except +
@@ -340,9 +342,10 @@ cdef extern from "api/cpp/cvc5.h" namespace "cvc5::api":
         bint operator==(const Term&) except +
         bint operator!=(const Term&) except +
         Term operator[](size_t idx) except +
+        uint64_t getId() except +
         Kind getKind() except +
         Sort getSort() except +
-        Term substitute(const vector[Term] es, const vector[Term] & reps) except +
+        Term substitute(const vector[Term] & es, const vector[Term] & reps) except +
         bint hasOp() except +
         Op getOp() except +
         bint isNull() except +
diff --git a/src/api/python/cvc5.pxi b/src/api/python/cvc5.pxi
index 11742ca08..f38a953ee 100644
--- a/src/api/python/cvc5.pxi
+++ b/src/api/python/cvc5.pxi
@@ -314,6 +314,9 @@ cdef class Op:
 
     def getKind(self):
         return kind(<int> self.cop.getKind())
+    
+    def isIndexed(self):
+        return self.cop.isIndexed()
 
     def isNull(self):
         return self.cop.isNull()
@@ -692,8 +695,11 @@ cdef class Solver:
 
     def mkInteger(self, val):
         cdef Term term = Term(self)
-        integer = int(val)
-        term.cterm = self.csolver.mkInteger("{}".format(integer).encode())
+        if isinstance(val, str):
+            term.cterm = self.csolver.mkInteger(<const string &> str(val).encode())
+        else:
+            assert(isinstance(val, int))
+            term.cterm = self.csolver.mkInteger((<int?> val))
         return term
 
     def mkReal(self, val, den=None):
@@ -1477,6 +1483,10 @@ cdef class Term:
     def __hash__(self):
         return ctermhash(self.cterm)
 
+
+    def getId(self):
+        return self.cterm.getId()
+
     def getKind(self):
         return kind(<int> self.cterm.getKind())
 
@@ -1485,20 +1495,33 @@ cdef class Term:
         sort.csort = self.cterm.getSort()
         return sort
 
-    def substitute(self, list es, list replacements):
+    def substitute(self, term_or_list_1, term_or_list_2):
+        # The resulting term after substitution
+        cdef Term term = Term(self.solver)
+        # lists for substitutions
         cdef vector[c_Term] ces
         cdef vector[c_Term] creplacements
-        cdef Term term = Term(self.solver)
-
-        if len(es) != len(replacements):
-            raise RuntimeError("Expecting list inputs to substitute to "
-                               "have the same length but got: "
-                               "{} and {}".format(len(es), len(replacements)))
-
-        for e, r in zip(es, replacements):
-            ces.push_back((<Term?> e).cterm)
-            creplacements.push_back((<Term?> r).cterm)
+        
+        # normalize the input parameters to be lists
+        if isinstance(term_or_list_1, list):
+            assert isinstance(term_or_list_2, list)
+            es = term_or_list_1
+            replacements = term_or_list_2
+            if len(es) != len(replacements):
+                raise RuntimeError("Expecting list inputs to substitute to "
+                                   "have the same length but got: "
+                                   "{} and {}".format(len(es), len(replacements)))
+
+            for e, r in zip(es, replacements):
+                ces.push_back((<Term?> e).cterm)
+                creplacements.push_back((<Term?> r).cterm)
 
+        else:
+            # add the single elements to the vectors
+            ces.push_back((<Term?> term_or_list_1).cterm)
+            creplacements.push_back((<Term?> term_or_list_2).cterm)
+        
+        # call the API substitute method with lists
         term.cterm = self.cterm.substitute(ces, creplacements)
         return term