improvements to python plugin

countess/plugins/python.py

@@ -1,26 +1,32 @@
 from types import CodeType
+import builtins
 
 import pandas as pd
 
 from countess import VERSION
 from countess.core.logger import Logger
 from countess.core.parameters import BooleanParam, PerColumnArrayParam, TextParam
-from countess.core.plugins import PandasTransformRowToDictPlugin
+from countess.core.plugins import PandasTransformDictToDictPlugin
 
 # XXX pretty sure this is a job for ast.parse rather than just
 # running compile() and exec() but that can wait.
-
+# Builtins are restricted but there's still plenty of things which
+# could go wrong here.
 
 # These types will get copied to columns, anything else
 # (eg: classes, methods, functions) won't.
 
 SIMPLE_TYPES = set((bool, int, float, str, tuple, list))
 
-# XXX should probably actually be based on
-# PandasTransformDictToDictPlugin
-# which is a bit more efficient.
+SAFE_BUILTINS = {
+    x: builtins.__dict__[x]
+    for x in
+        "abs all any ascii bin bool bytearray bytes chr complex dict divmod enumerate "
+        "filter float format frozenset hash hex id int len list map max min oct "
+        "ord pow range reversed round set slice sorted str sum tuple type zip".split()
+}
 
-class PythonPlugin(PandasTransformRowToDictPlugin):
+class PythonPlugin(PandasTransformDictToDictPlugin):
     name = "Python Code"
     description = "Apply python code to each row."
     additional = """
@@ -38,21 +44,21 @@ class PythonPlugin(PandasTransformRowToDictPlugin):
 
     code_object = None
 
-    def process_row(self, row: pd.Series, logger: Logger):
+
+    def process_dict(self, data: dict, logger: Logger):
         assert isinstance(self.parameters["code"], TextParam)
         assert isinstance(self.parameters["columns"], PerColumnArrayParam)
         assert isinstance(self.code_object, CodeType)
 
-        row_dict = dict(row)
         try:
-            exec(self.code_object, {}, row_dict)  # pylint: disable=exec-used
+            exec(self.code_object, {"__builtins__": SAFE_BUILTINS}, data)  # pylint: disable=exec-used
         except Exception as exc:  # pylint: disable=broad-exception-caught
             logger.exception(exc)
 
         column_parameters = list(zip(self.input_columns, self.parameters["columns"].params))
         columns_to_remove = set(col for col, param in column_parameters if not param.value)
 
-        return dict((k, v) for k, v in row_dict.items() if k not in columns_to_remove and type(v) in SIMPLE_TYPES)
+        return dict((k, v) for k, v in data.items() if k not in columns_to_remove and type(v) in SIMPLE_TYPES)
 
     def process_dataframe(self, dataframe: pd.DataFrame, logger: Logger) -> pd.DataFrame:
         """Override parent class because we a) want to reset