From 54af29dbdb6c266bae2590e68d4a0b2c5e61279c Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 24 Dec 2025 23:57:10 -0500 Subject: [PATCH] trange can just be a function (#13827) --- test/unit/test_tqdm.py | 4 ++-- tinygrad/helpers.py | 3 +-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/test/unit/test_tqdm.py b/test/unit/test_tqdm.py index 52c64c7ea6..2ba3f2fe4a 100644 --- a/test/unit/test_tqdm.py +++ b/test/unit/test_tqdm.py @@ -3,7 +3,7 @@ from unittest.mock import patch from io import StringIO from collections import namedtuple from tqdm import tqdm -from tinygrad.helpers import tqdm as tinytqdm, trange as tinytrange +from tinygrad.helpers import tqdm as tinytqdm, trange import numpy as np def _get_iter_per_second(raw:str) -> float: @@ -165,7 +165,7 @@ class TestProgressBar(unittest.TestCase): mock_stderr.truncate(0) # compare bars at each iteration (only when tinytqdm bar has been updated) - for n in (bar := tinytrange(total, desc="Test")): + for n in (bar := trange(total, desc="Test")): if bar.i % bar.skip != 0: continue tiny_output = mock_stderr.getvalue().split("\r")[-1].rstrip() iters_per_sec = float(tiny_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0 diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 8069ae842f..e0000f1590 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -523,8 +523,7 @@ class tqdm(Generic[T]): @classmethod def write(cls, s:str): print(f"\r\033[K{s}", flush=True, file=sys.stderr) -class trange(tqdm): - def __init__(self, n:int, **kwargs): super().__init__(iterable=range(n), total=n, **kwargs) +def trange(n:int, **kwargs) -> tqdm[int]: return tqdm(range(n), total=n, **kwargs) class disable_gc(contextlib.ContextDecorator): def __enter__(self):