refactor: move start/end trim to utils.read

src/stepcount/stepcount.py

@@ -82,18 +82,14 @@ def main():
     data, info_read = utils.read(
         args.filepath, 
         usecols=args.txyz, 
-        resample_hz=30 if args.model_type == 'ssl' else None,
+        start_time=args.start,
+        end_time=args.end,
         sample_rate=args.sample_rate, 
+        resample_hz=30 if args.model_type == 'ssl' else None,
         verbose=verbose
     )
     info.update(info_read)
 
-    # Set start/end times, if given
-    if args.start is not None:
-        data = data.loc[args.start:]
-    if args.end is not None:
-        data = data.loc[:args.end]
-
     # Exclusion: first/last days
     if args.exclude_first_last is not None:
         data = utils.exclude_first_last_days(data, args.exclude_first_last)

src/stepcount/utils.py

@@ -12,8 +12,10 @@
 def read(
     filepath: str,
     usecols: str = 'time,x,y,z',
-    resample_hz: str = 'uniform',
+    start_time: str = None,
+    end_time: str = None,
     sample_rate: float = None,
+    resample_hz: str = 'uniform',
     verbose: bool = True
 ):
     """
@@ -91,8 +93,6 @@ def read(
             "Device": ftype,
             "Filesize(MB)": fsize,
             "SampleRate": sample_rate,
-            "StartTime": data.index[0].strftime('%Y-%m-%d %H:%M:%S'),
-            "EndTime": data.index[-1].strftime('%Y-%m-%d %H:%M:%S')
         })
 
     elif ftype in (".cwa", ".gt3x", ".bin"):
@@ -112,6 +112,18 @@ def read(
     if 'ResampleRate' not in info:
         info['ResampleRate'] = info['SampleRate']
 
+    # Trim the data if start/end times are specified
+    if start_time is not None:
+        data = data.loc[start_time:]
+    if end_time is not None:
+        data = data.loc[:end_time]
+
+    # Update start/end times in metadata
+    info.update({
+        "StartTime": data.index[0].strftime('%Y-%m-%d %H:%M:%S'),
+        "EndTime": data.index[-1].strftime('%Y-%m-%d %H:%M:%S')
+    })
+
     return data, info