feat: Rename command + test #5 #4

i2mint · Jan 23, 2024 · f2dd1d9 · f2dd1d9
1 parent 22b1a70
commit f2dd1d9
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 15 deletions.
diff --git a/tabled/multi.py b/tabled/multi.py
@@ -2,7 +2,7 @@
 
 from functools import cached_property
 import pandas as pd
-from typing import TypeVar, Callable, KT, Union, Iterable, Mapping
+from typing import TypeVar, Callable, KT, Union, Iterable, Mapping, Dict
 
 Column = TypeVar('Column')
 TableKey = TypeVar('TableKey')
@@ -78,6 +78,10 @@ class Join:
 class Remove:
     fields: Union[str, Iterable[str]]
 
+@dataclass
+class Rename:
+    rename_mapping: Dict[str, str]
+
 
 def set_scope_value(scope, key, value):
     scope[key] = value
@@ -97,10 +101,15 @@ def remove_func(scope, command):
     scope['cumul'] = scope['cumul'].drop(columns=command.fields)
 
 
+def rename_func(scope, command):
+    for old_col, new_col in command.rename_mapping.items():
+        scope['cumul'] = scope['cumul'].rename(columns={old_col: new_col})
+
 dflt_tables_interpreter_map = {
     Load: load_func,
     Join: join_func,
     Remove: remove_func,
+    Rename: rename_func,
 }
 
 

diff --git a/tabled/tests/join_tables.py b/tabled/tests/join_tables.py
@@ -132,7 +132,7 @@ def test_compute_join_resolution(tables):
 
 
 def test_execute_commands_simply():
-    from tabled.multi import Join, Remove, Load
+    from tabled.multi import Join, Remove, Load, Rename
 
     # ---------------------------------------------
     # First silly test
@@ -156,7 +156,7 @@ def test_execute_commands_simply():
 
     tables = {'table1': table1, 'table2': table2, 'table3': table3}
 
-    commands = [Load('table1'), Remove(['Name']), Join('table3')]
+    commands = [Load('table1'), Rename({'Name': 'First Name'}), Remove(['First Name']), Join('table3')]
 
     scope = tables
     extra_scope = dict()
@@ -172,6 +172,8 @@ def are_equal(a, b):
     next(it)
     assert are_equal(extra_scope['cumul'], scope['table1'])
     next(it)
+    assert list(extra_scope['cumul']) == ['ID', 'First Name']
+    next(it)
     assert are_equal(extra_scope['cumul'], pd.DataFrame({'ID': [1, 2, 3]}))
     next(it)
     assert list(extra_scope['cumul']) == ['ID', 'Salary']
@@ -199,17 +201,17 @@ def extracted_dataframes():
     return dfs_aeroports_frequentes, dfs_aeroports_vastes
 
 
-def test_execute_commands_wiki(extracted_dataframes):
-    from tabled.multi import Join, Remove, Load
-    table1_wiki = extracted_dataframes[0]
-    table2_wiki =  extracted_dataframes[1]
+# def test_execute_commands_wiki(extracted_dataframes):
+#     from tabled.multi import Join, Remove, Load
+#     table1_wiki = extracted_dataframes[0]
+#     table2_wiki =  extracted_dataframes[1]
 
-    tables = {'table1_wiki': table1_wiki, 'table2_wiki': table2_wiki}
-    commands = [Load('table1_wiki'), Join('table2_wiki')]
+#     tables = {'table1_wiki': table1_wiki, 'table2_wiki': table2_wiki}
+#     commands = [Load('table1_wiki'), Join('table2_wiki')]
 
-    scope = tables
-    extra_scope = dict()
-    it = execute_table_commands(commands, scope, extra_scope=extra_scope)
-    next(it)
-    next(it)
-    assert extra_scope['cumul'].shape[0] == 1
+#     scope = tables
+#     extra_scope = dict()
+#     it = execute_table_commands(commands, scope, extra_scope=extra_scope)
+#     next(it)
+#     next(it)
+#     assert extra_scope['cumul'].shape[0] == 1