+ linter

Author: teodor-r

rectools/model_selection/last_n_split.py

@@ -17,6 +17,7 @@
 
 import numpy as np
 import pandas as pd
+
 from rectools import Columns
 from rectools.dataset import Interactions
 from rectools.model_selection.splitter import Splitter
@@ -102,11 +103,7 @@ def _split_without_filter(
         idx = pd.RangeIndex(0, len(df))
 
         # Here we guarantee that last appeared interaction in df will have lowest rank when datetime is not unique
-        time_order = (
-            df.groupby(Columns.User)[Columns.Datetime]
-            .rank(method="first", ascending=True)
-            .astype(int)
-        )
+        time_order = df.groupby(Columns.User)[Columns.Datetime].rank(method="first", ascending=True).astype(int)
         n_interactions = df.groupby(Columns.User).transform("size").astype(int)
         inv_ranks = n_interactions - time_order + 1
 

rectools/models/nn/item_net.py

@@ -486,4 +486,4 @@ def forward(self, items: torch.Tensor) -> torch.Tensor:
     @property
     def out_dim(self) -> int:
         """Return item net constructor output dimension."""
-        return self.item_net_blocks[0].out_dim  # type: ignore[return-value]
+        return self.item_net_blocks[0].out_dim

tests/model_selection/test_last_n_split.py

@@ -40,6 +40,7 @@ def _shuffle(values: tp.Sequence[int]) -> tp.List[int]:
             return sorted(inv_shuffle_arr[values])
 
         return _shuffle
+
     @pytest.fixture
     def interactions_equal_timestamps(self, shuffle_arr: np.ndarray) -> Interactions:
         df = pd.DataFrame(
@@ -62,37 +63,29 @@ def interactions_equal_timestamps(self, shuffle_arr: np.ndarray) -> Interactions
     @pytest.mark.parametrize(
         "swap_targets,expected_test_ids, target_item",
         (
-            (
-                False,
-                {9, 7, 8},
-                6
-            ),
-            (
-                True,
-                {9, 7, 8},
-                3
-            ),
+            (False, {9, 7, 8}, 6),
+            (True, {9, 7, 8}, 3),
         ),
     )
     def test_correct_last_interactions(
         self,
         interactions_equal_timestamps: Interactions,
         swap_targets: bool,
-        expected_test_ids: tp.List[int],
+        expected_test_ids: tp.Set[int],
         target_item: int,
     ) -> None:
         # Do not using shuffle fixture, otherwise no valid answers
         interactions_et = interactions_equal_timestamps
         splitter = LastNSplitter(1, 1, False, False, False)
         if swap_targets:
             df_swap = interactions_equal_timestamps.df
-            df_swap.iloc[[4,9]] = df_swap.iloc[[9,4]]
+            df_swap.iloc[[4, 9]] = df_swap.iloc[[9, 4]]
             interactions_et = Interactions(df_swap)
         loo_split = list(splitter.split(interactions_et, collect_fold_stats=True))
         target_ids = loo_split[0][1]
         assert set(target_ids) == expected_test_ids
-        assert set(loo_split[0][0]) ==  set(range(len(interactions_et.df))) - expected_test_ids
-        assert target_item in  set(interactions_et.df.iloc[target_ids][Columns.Item])
+        assert set(loo_split[0][0]) == set(range(len(interactions_et.df))) - expected_test_ids
+        assert target_item in set(interactions_et.df.iloc[target_ids][Columns.Item])
 
     @pytest.fixture
     def interactions(self, shuffle_arr: np.ndarray) -> Interactions: