♻️ Refactor code

src/viadot/orchestration/prefect/flows/tm1_to_parquet.py

@@ -14,7 +14,7 @@
     retries=1,
     retry_delay_seconds=60,
 )
-def tm1_to_parquet( # noqa: PLR0913
+def tm1_to_parquet(  # noqa: PLR0913
     path: str | None = None,
     mdx_query: str | None = None,
     cube: str | None = None,

src/viadot/sources/__init__.py

@@ -49,6 +49,7 @@
     "Sftp",
     "Sharepoint",
     "Supermetrics",
+    "TM1",
     "Trino",
     "UKCarbonIntensity",
     "VidClub",

src/viadot/sources/tm1.py

@@ -193,7 +193,7 @@ def to_df(self, if_empty: Literal["warn", "fail", "skip"] = "skip") -> pd.DataFr
         if self.mdx_query is not None and (
             self.cube is not None or self.view is not None
         ):
-            error_msg="Specify only one: MDX query or cube and view."
+            error_msg = "Specify only one: MDX query or cube and view."
             raise ValidationError(error_msg)
         if self.cube is not None and self.view is not None:
             df = conn.cubes.cells.execute_view_dataframe(

tests/unit/test_tm1.py

@@ -8,18 +8,15 @@
 from viadot.sources.tm1 import TM1Credentials
 
 
-data = {
-  "a": [420, 380, 390],
-  "b": [50, 40, 45]
-}
+data = {"a": [420, 380, 390], "b": [50, 40, 45]}
 
 @pytest.fixture
 def tm1_credentials():
     return TM1Credentials(
         username="test_user",
         password="test_password",  # pragma: allowlist secret # noqa: S106
-        address = "localhost",
-        port = "12356",
+        address="localhost",
+        port="12356",
     )
 
 @pytest.fixture
@@ -31,11 +28,11 @@ def tm1(tm1_credentials: TM1Credentials):
             "address": tm1_credentials.address,
             "port": tm1_credentials.port,
         },
-        verify = False,
-        cube = "test_cube",
-        view = "test_view",
-        dimension = "test_dim",
-        hierarchy = "test_hier",
+        verify=False,
+        cube="test_cube",
+        view="test_view",
+        dimension="test_dim",
+        hierarchy="test_hier",
     )
 
 
@@ -64,35 +61,35 @@ def test_tm1_initialization(tm1_mock):
     )
 
 def test_get_cubes_names(tm1):
-    with patch('viadot.sources.tm1.TM1Service') as mock:
+    with patch("viadot.sources.tm1.TM1Service") as mock:
         instance = mock.return_value
-        instance.cubes.get_all_names.return_value = ['msg', 'cuve']
+        instance.cubes.get_all_names.return_value = ["msg", "cuve"]
         result = tm1.get_cubes_names()
-        assert result == ['msg', 'cuve']
+        assert result == ["msg", "cuve"]
 
 def test_get_views_names(tm1):
-    with patch('viadot.sources.tm1.TM1Service') as mock:
+    with patch("viadot.sources.tm1.TM1Service") as mock:
         instance = mock.return_value
-        instance.views.get_all_names.return_value = ['view1', 'view1']
+        instance.views.get_all_names.return_value = ["view1", "view1"]
         result = tm1.get_views_names()
-        assert result == ['view1', 'view1']
+        assert result == ["view1", "view1"]
 
 def test_get_dimensions_names(tm1):
-    with patch('viadot.sources.tm1.TM1Service') as mock:
+    with patch("viadot.sources.tm1.TM1Service") as mock:
         instance = mock.return_value
-        instance.dimensions.get_all_names.return_value = ['dim1', 'dim2']
+        instance.dimensions.get_all_names.return_value = ["dim1", "dim2"]
         result = tm1.get_dimensions_names()
-        assert result == ['dim1', 'dim2']
+        assert result == ["dim1", "dim2"]
 
 def test_get_available_elements(tm1):
-    with patch('viadot.sources.tm1.TM1Service') as mock:
+    with patch("viadot.sources.tm1.TM1Service") as mock:
         instance = mock.return_value
-        instance.elements.get_element_names.return_value = ['el1', 'el2']
+        instance.elements.get_element_names.return_value = ["el1", "el2"]
         result = tm1.get_available_elements()
-        assert result == ['el1', 'el2']
+        assert result == ["el1", "el2"]
 
 def test_to_df(tm1):
-    with patch('viadot.sources.tm1.TM1Service') as mock:
+    with patch("viadot.sources.tm1.TM1Service") as mock:
         instance = mock.return_value
         instance.cubes.cells.execute_view_dataframe.return_value = pd.DataFrame(data)
         result = tm1.to_df()