Fixed evaluation environments

jpmml · Oct 28, 2024 · 8ca1d9f · 8ca1d9f
1 parent 0d5c489
commit 8ca1d9f
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/sklearn2pmml/preprocessing/tests/__init__.py b/sklearn2pmml/preprocessing/tests/__init__.py
@@ -405,6 +405,13 @@ def test_func_transform(self):
 		X = numpy.array([[1.5], [0.0], [-3.0]])
 		self.assertEqual([[1], [0], [-1]], transformer.fit_transform(X).tolist())
 
+	def test_regex_transform(self):
+		transformer = ExpressionTransformer("pcre2.substitute(r'B+', r'c', X[0])")
+		X = numpy.array([["ABBA"]])
+		self.assertEqual([["AcA"]], transformer.fit_transform(X).tolist())
+		transformer = ExpressionTransformer("re.sub(r'B+', r'c', X[0])")
+		self.assertEqual([["AcA"]], transformer.fit_transform(X).tolist())
+
 	def test_sequence_transform(self):
 		X = DataFrame([[None], [1], [None]], columns = ["a"])
 		mapper = DataFrameMapper([

diff --git a/sklearn2pmml/util/__init__.py b/sklearn2pmml/util/__init__.py
@@ -181,7 +181,7 @@ def to_expr(expr):
 	else:
 		raise TypeError()
 
-def to_expr_func(expr, modules = ["math", "re", "pcre", "numpy", "pandas", "scipy"]):
+def to_expr_func(expr, modules = ["math", "re", "pcre", "pcre2", "numpy", "pandas", "scipy"]):
 	env = dict()
 
 	if isinstance(expr, str):