diff --git a/api/onnx_web/worker/command.py b/api/onnx_web/worker/command.py index fcbc5b7b..40f4ed77 100644 --- a/api/onnx_web/worker/command.py +++ b/api/onnx_web/worker/command.py @@ -51,13 +51,16 @@ class Progress: "total": self.total, } - def complete(self) -> bool: + def is_complete(self) -> bool: return self.current >= self.total - def empty(self) -> bool: + def is_empty(self) -> bool: # TODO: what if total is also 0? return self.current == 0 + def update(self, current: int) -> "Progress": + return Progress(current, self.total) + class ProgressCommand: device: str diff --git a/api/onnx_web/worker/context.py b/api/onnx_web/worker/context.py index 1f5c57a0..ebfbd462 100644 --- a/api/onnx_web/worker/context.py +++ b/api/onnx_web/worker/context.py @@ -140,13 +140,13 @@ class WorkerContext: raise CancelledException("job has been cancelled") # update current progress counters - self.steps.current = steps + self.steps = self.steps.update(steps) if stages is not None: - self.stages.current = stages + self.stages = self.stages.update(stages) if tiles is not None: - self.tiles.current = tiles + self.tiles = self.tiles.update(tiles) # TODO: result should really be part of context at this point result = None @@ -174,19 +174,19 @@ class WorkerContext: if total > 0: self.steps = Progress(current, total) else: - self.steps.current = current + self.steps = self.steps.update(current) def set_stages(self, current: int, total: int = 0) -> None: if total > 0: self.stages = Progress(current, total) else: - self.stages.current = current + self.stages = self.stages.update(current) def set_tiles(self, current: int, total: int = 0) -> None: if total > 0: self.tiles = Progress(current, total) else: - self.tiles.current = current + self.tiles = self.tiles.update(current) def set_totals(self, steps: int, stages: int = 0, tiles: int = 0) -> None: self.steps.total = steps