diff --git a/python/taichi/_main.py b/python/taichi/_main.py index ddd9446cbf04a..2b56c3334a239 100644 --- a/python/taichi/_main.py +++ b/python/taichi/_main.py @@ -225,14 +225,14 @@ def on_mouse_click_callback(example_name): @register def example(self, arguments: list = sys.argv[2:]): """Run an example by name (or name.py)""" - def colormap(name): + def colormap(index, name): from colorsys import hls_to_rgb # pylint: disable=C0415 x = (ord(name[0].upper()) - 64.0) / 26.0 r, g, b = hls_to_rgb(x, 0.4, 1.0) r = hex(int(r * 255) % 16)[2:] g = hex(int(g * 255) % 16)[2:] b = hex(int(b * 255) % 16)[2:] - return f"[#{r}{r}{g}{g}{b}{b}]{name}" + return f"{index}: [#{r}{r}{g}{g}{b}{b}]{name}" console = Console() table = rich.table.Table( @@ -249,7 +249,8 @@ def colormap(name): names = sorted(choices.keys()) for k in range(nrows): table.add_row( - *[colormap(names[j]) for j in range(k, len(choices), nrows)]) + * + [colormap(j, names[j]) for j in range(k, len(choices), nrows)]) parser = argparse.ArgumentParser(prog='ti example', description=f"{self.example.__doc__}") @@ -258,6 +259,8 @@ def colormap(name): choices.keys()), choices=sorted(choices.keys()), help=console.print(table), + nargs="?", + default=None, metavar="name") parser.add_argument( '-p', @@ -285,17 +288,26 @@ def colormap(name): args = parser.parse_args(arguments) examples_dir = TaichiMain._get_examples_dir() - target = str( - (examples_dir / choices[args.name] / f"{args.name}.py").resolve()) + example_name = args.name + if example_name is None: + index = int(input("Please input the number of the example: ")) + while not 0 <= index < len(names): + index = int( + input( + f"Example [{index}] does not exist. Please try again: " + )) + example_name = names[index] + target = str((examples_dir / choices[example_name] / + f"{example_name}.py").resolve()) # path for examples needs to be modified for implicit relative imports - sys.path.append(str((examples_dir / choices[args.name]).resolve())) + sys.path.append(str((examples_dir / choices[example_name]).resolve())) # Short circuit for testing if self.test_mode: return args if args.save: - print(f"Saving example {args.name} to current directory...") + print(f"Saving example {example_name} to current directory...") shutil.copy(target, '.') return 0 @@ -310,12 +322,10 @@ def colormap(name): print(f.read()) return 0 - print(f"Running example {args.name} ...") + print(f"Running example {example_name} ...") runpy.run_path(target, run_name='__main__') - return None - @staticmethod @register def changelog(arguments: list = sys.argv[2:]): diff --git a/python/taichi/examples/simulation/vortex_rings.py b/python/taichi/examples/simulation/vortex_rings.py index 3d75655de18da..01453a947f03c 100644 --- a/python/taichi/examples/simulation/vortex_rings.py +++ b/python/taichi/examples/simulation/vortex_rings.py @@ -79,7 +79,7 @@ def main(): init_tracers() gui = ti.GUI("Vortex Rings", (1024, 512), background_color=0xFFFFFF) - for T in range(1000): + while gui.running: for i in range(4): # substeps advect() integrate_vortex()