Skip to content

Commit

Permalink
Dataclass instead of tuple
Browse files Browse the repository at this point in the history
  • Loading branch information
ecomodeller committed Nov 21, 2024
1 parent b441e98 commit ffcc36d
Showing 1 changed file with 27 additions and 14 deletions.
41 changes: 27 additions & 14 deletions mikeio/generic.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
"""Generic functions for working with all types of dfs files."""

from __future__ import annotations
from dataclasses import dataclass
import math
import os
import pathlib
Expand Down Expand Up @@ -597,10 +598,7 @@ def extract(

is_layered_dfsu = dfs_i.ItemInfo[0].Name == "Z coordinate"

file_start_new, start_step, start_sec, end_step, end_sec = _parse_start_end(
dfs_i.FileInfo.TimeAxis, start, end
)
timestep = _parse_step(dfs_i.FileInfo.TimeAxis, step)
time = _parse_time(dfs_i.FileInfo.TimeAxis, start, end, step)
item_numbers = _valid_item_numbers(
dfs_i.ItemInfo, items, ignore_first=is_layered_dfsu
)
Expand All @@ -612,28 +610,28 @@ def extract(
dfs_o = _clone(
str(infilename),
str(outfilename),
start_time=file_start_new,
timestep=timestep,
start_time=time.file_start_new,
timestep=time.timestep,
items=item_numbers,
)

file_start_shift = 0
if file_start_new is not None:
if time.file_start_new is not None:
file_start_orig = dfs_i.FileInfo.TimeAxis.StartDateTime
file_start_shift = (file_start_new - file_start_orig).total_seconds()
file_start_shift = (time.file_start_new - file_start_orig).total_seconds()

timestep_out = -1
for timestep in range(start_step, end_step, step):
for timestep in range(time.start_step, time.end_step, step):
for item_out, item in enumerate(item_numbers):
itemdata = dfs_i.ReadItemTimeStep((item + 1), timestep)
time_sec = itemdata.Time

if time_sec > end_sec:
if time_sec > time.end_sec:
dfs_i.Close()
dfs_o.Close()
return

if time_sec >= start_sec:
if time_sec >= time.start_sec:
if item == item_numbers[0]:
timestep_out = timestep_out + 1
time_sec_out = time_sec - file_start_shift
Expand All @@ -647,11 +645,24 @@ def extract(
dfs_o.Close()


def _parse_start_end(
@dataclass
class TimeInfo:
"Parsed time information."

file_start_new: datetime | None
start_step: int
start_sec: float
end_step: int
end_sec: float
timestep: float | None


def _parse_time(
time_axis: TimeAxis,
start: int | float | str | datetime,
end: int | float | str | datetime,
) -> tuple[datetime | None, int, float, int, float]: # TODO better return type
step: int,
) -> TimeInfo:
"""Helper function for parsing start and end arguments."""
n_time_steps = time_axis.NumberOfTimeSteps
file_start_datetime = time_axis.StartDateTime
Expand Down Expand Up @@ -727,7 +738,9 @@ def _parse_start_end(
if start_sec > file_start_sec:
file_start_new = file_start_datetime + timedelta(seconds=start_sec)

return file_start_new, start_step, start_sec, end_step, end_sec
timestep = _parse_step(time_axis, step)

return TimeInfo(file_start_new, start_step, start_sec, end_step, end_sec, timestep)


def _parse_step(time_axis: TimeAxis, step: int) -> float | None:
Expand Down

0 comments on commit ffcc36d

Please sign in to comment.