DataFrame column tool

countess/plugins/column.py

@@ -0,0 +1,62 @@
+import pandas as pd
+
+from countess import VERSION
+from countess.core.logger import Logger
+from countess.core.parameters import (
+    BooleanParam,
+    DataTypeOrNoneChoiceParam,
+    PerColumnArrayParam,
+    StringParam,
+    TabularMultiParam,
+)
+from countess.core.plugins import PandasSimplePlugin
+
+
+class ColumnToolPlugin(PandasSimplePlugin):
+    name = "DataFrame Column Tool"
+    description = "Alter Columns of a DataFrame"
+    version = VERSION
+
+    parameters = {
+        "columns": PerColumnArrayParam(
+            "Columns",
+            TabularMultiParam(
+                "Column",
+                {
+                    "rename": StringParam("Name"),
+                    "datatype": DataTypeOrNoneChoiceParam("Column Type"),
+                    "index": BooleanParam("Index?"),
+                },
+            ),
+        )
+    }
+
+    def process_dataframe(self, dataframe: pd.DataFrame, logger: Logger) -> pd.DataFrame:
+        column_parameters = list(zip(self.input_columns, self.parameters["columns"]))
+
+        drop_columns = [column_name for column_name, parameter in column_parameters if parameter.datatype.is_none()]
+
+        type_columns = {
+            column_name: parameter.datatype.get_selected_type()
+            for column_name, parameter in column_parameters
+            if parameter.datatype.is_not_none()
+        }
+
+        index_columns = [
+            column_name
+            for column_name, parameter in column_parameters
+            if parameter.index.value and parameter.datatype.is_not_none()
+        ]
+
+        rename_columns = {
+            column_name: parameter.rename.value.strip()
+            for column_name, parameter in column_parameters
+            if parameter.rename.value and parameter.rename.value.strip() and parameter.datatype.is_not_none()
+        }
+
+        dataframe = dataframe.reset_index().drop(columns=drop_columns).astype(type_columns)
+
+        if index_columns:
+            dataframe = dataframe.set_index(index_columns)
+
+        return dataframe.rename(columns=rename_columns)

pyproject.toml

@@ -67,6 +67,7 @@ mutagenize = "countess.plugins.mutagenize:MutagenizePlugin"
 data_table = "countess.plugins.data_table:DataTablePlugin"
 correlation = "countess.plugins.correlation:CorrelationPlugin"
 hgvs_parser = "countess.plugins.hgvs_parser:HgvsParserPlugin"
+column = "countess.plugins.column:ColumnToolPlugin"
 
 [project.entry-points.gui_scripts]
 countess_gui = "countess.gui.main:main"