Use textwrap to shorten debug description

Author: samedii

datastream/dataset.py

@@ -4,6 +4,8 @@
     Tuple, Callable, Any, Union, List, TypeVar, Generic, Dict, Optional
 )
 from pathlib import Path
+from functools import lru_cache
+import textwrap
 import inspect
 import numpy as np
 import pandas as pd
@@ -132,9 +134,7 @@ def composed_fn(dataframe, index):
             try:
                 return function(item)
             except Exception as e:
-                item_text = str(item)
-                if len(item_text) > 79:
-                    item_text = item_text[:79] + ' ...'
+                item_text = textwrap.shorten(str(item), width=79)
 
                 raise Exception('\n'.join([
                     repr(e),
@@ -487,10 +487,11 @@ def zip(datasets: List[Dataset]) -> Dataset[Tuple]:
             ))
         )
 
-    def cache(self, key_column):
-        '''Cache dataset in-memory based on key column.'''
-        from functools import lru_cache
-
+    def cache(
+        self,
+        key_column: str,
+    ):
+        '''Cache intermediate step in-memory based on key column.'''
         key_mapping = dict(zip(
             self.dataframe[key_column],
             range(len(self)),

datastream/datastream.py

@@ -12,6 +12,7 @@
 )
 import numpy as np
 import torch
+from pathlib import Path
 
 from datastream import Dataset
 from datastream.samplers import (
@@ -253,7 +254,10 @@ def multi_sample(self: Datastream[T], n: int) -> Datastream[T]:
             MultiSampler.from_number(n, self.dataset),
         )
 
-    def cache(self, key_column):
+    def cache(
+        self,
+        key_column: str,
+    ):
         '''Cache dataset in-memory. See :func:`Dataset.cache` for details.'''
         return Datastream(
             self.dataset.cache(key_column),