refactor: add type casting for return values in read_transactions_from_spm

Author: jacksonpradolima

gsppy/cli.py

@@ -216,7 +216,7 @@ def read_transactions_from_spm(file_path: str) -> List[List[str]]:
     try:
         from gsppy.utils import read_transactions_from_spm as read_spm
 
-        return read_spm(file_path, return_mappings=False)
+        return cast(List[List[str]], read_spm(file_path, return_mappings=False))
     except Exception as e:
         msg = f"Error reading transaction data from SPM file '{file_path}': {e}"
         logging.error(msg)
@@ -363,15 +363,15 @@ def _load_dataframe_format(
 ) -> Union[List[List[str]], List[List[Tuple[str, float]]]]:
     """
     Load transactions from DataFrame formats (Parquet/Arrow).
-    
+
     Parameters:
         file_path: Path to the file
         file_extension: File extension (lowercase)
         transaction_col: Transaction ID column name
         item_col: Item column name
         timestamp_col: Timestamp column name
         sequence_col: Sequence column name
-        
+
     Returns:
         Loaded transactions
     """
@@ -405,7 +405,7 @@ def _load_transactions_by_format(
 ) -> Union[List[List[str]], List[List[Tuple[str, float]]]]:
     """
     Load transactions based on specified format.
-    
+
     Parameters:
         file_path: Path to the file
         file_format: Format string (lowercase)
@@ -415,10 +415,10 @@ def _load_transactions_by_format(
         item_col: Item column name
         timestamp_col: Timestamp column name
         sequence_col: Sequence column name
-        
+
     Returns:
         Loaded transactions
-        
+
     Raises:
         ValueError: If format is unknown
     """
@@ -429,9 +429,7 @@ def _load_transactions_by_format(
     elif file_format == FileFormat.CSV.value:
         return read_transactions_from_csv(file_path)
     elif file_format in (FileFormat.PARQUET.value, FileFormat.ARROW.value):
-        return _load_dataframe_format(
-            file_path, file_extension, transaction_col, item_col, timestamp_col, sequence_col
-        )
+        return _load_dataframe_format(file_path, file_extension, transaction_col, item_col, timestamp_col, sequence_col)
     elif file_format == FileFormat.AUTO.value:
         # Auto-detect format
         if is_dataframe_format:

tests/test_spm_format.py

@@ -10,7 +10,7 @@
 
 import os
 import tempfile
-from typing import Generator
+from typing import Generator, cast
 
 import pytest
 
@@ -110,7 +110,8 @@ def test_basic_parsing(self, simple_spm_file: str):
 
     def test_basic_parsing_with_mappings(self, simple_spm_file: str):
         """Test SPM parsing with token mappings."""
-        transactions, str_to_int, int_to_str = read_transactions_from_spm(simple_spm_file, return_mappings=True)
+        result = cast(tuple, read_transactions_from_spm(simple_spm_file, return_mappings=True))
+        transactions, str_to_int, int_to_str = result
 
         assert len(transactions) == 3
         assert len(str_to_int) == 6  # Unique tokens: 1, 2, 3, 4, 5, 6
@@ -261,7 +262,7 @@ def complex_spm_file(self) -> Generator[str, None, None]:
 
     def test_complex_parsing(self, complex_spm_file: str):
         """Test parsing complex SPM file."""
-        transactions = read_transactions_from_spm(complex_spm_file)
+        transactions = cast(list, read_transactions_from_spm(complex_spm_file))
 
         assert len(transactions) == 4
         assert transactions[0] == ["1", "2", "3", "1", "4", "5"]
@@ -271,7 +272,8 @@ def test_complex_parsing(self, complex_spm_file: str):
 
     def test_complex_with_mappings(self, complex_spm_file: str):
         """Test complex file with mappings."""
-        transactions, str_to_int, _ = read_transactions_from_spm(complex_spm_file, return_mappings=True)
+        result = cast(tuple, read_transactions_from_spm(complex_spm_file, return_mappings=True))
+        transactions, str_to_int, _ = result
 
         assert len(transactions) == 4
         # Unique tokens: 1, 2, 3, 4, 5