Skip to content

Commit 9ff0bba

Browse files
authored
Merge pull request #500 from willcl-ark/inquirer-quickstart
Use inquirer for quickstart
2 parents 6711ad8 + aab9443 commit 9ff0bba

File tree

1 file changed

+141
-59
lines changed

1 file changed

+141
-59
lines changed

src/warnet/main.py

Lines changed: 141 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -47,85 +47,162 @@ def cli():
4747
def quickstart():
4848
"""Setup warnet"""
4949
try:
50-
process = subprocess.Popen(
50+
# Requirements checks
51+
with subprocess.Popen(
5152
["/bin/bash", str(QUICK_START_PATH)],
5253
stdout=subprocess.PIPE,
5354
stderr=subprocess.STDOUT,
5455
universal_newlines=True,
5556
env=dict(os.environ, TERM="xterm-256color"),
56-
)
57-
for line in iter(process.stdout.readline, ""):
58-
click.echo(line, nl=False)
59-
process.stdout.close()
60-
return_code = process.wait()
61-
if return_code != 0:
62-
click.secho(
63-
f"Quick start script failed with return code {return_code}", fg="red", bold=True
64-
)
65-
click.secho("Install missing requirements before proceeding", fg="yellow")
57+
) as process:
58+
if process.stdout:
59+
for line in iter(process.stdout.readline, ""):
60+
click.echo(line, nl=False)
61+
return_code = process.wait()
62+
if return_code != 0:
63+
click.secho(
64+
f"Quick start script failed with return code {return_code}", fg="red", bold=True
65+
)
66+
click.secho("Install missing requirements before proceeding", fg="yellow")
67+
return False
68+
69+
# New project setup
70+
questions = [
71+
inquirer.Confirm(
72+
"create_project",
73+
message=click.style("Do you want to create a new project?", fg="blue", bold=True),
74+
default=True,
75+
),
76+
]
77+
answers = inquirer.prompt(questions)
78+
if answers is None:
79+
click.secho("Setup cancelled by user.", fg="yellow")
6680
return False
67-
68-
create_project = click.confirm(
69-
click.style("\nDo you want to create a new project?", fg="blue", bold=True),
70-
default=True,
71-
)
72-
if not create_project:
81+
if not answers["create_project"]:
7382
click.secho("\nSetup completed successfully!", fg="green", bold=True)
7483
return True
7584

76-
default_path = os.path.abspath(os.getcwd())
77-
project_path = click.prompt(
78-
click.style("\nEnter the project directory path", fg="blue", bold=True),
79-
default=default_path,
80-
type=click.Path(file_okay=False, dir_okay=True, resolve_path=True),
81-
)
82-
83-
custom_network = click.confirm(
84-
click.style("\nDo you want to create a custom network?", fg="blue", bold=True),
85-
default=True,
86-
)
87-
if not custom_network:
88-
create_warnet_project(Path(project_path))
85+
# Custom project setup
86+
questions = [
87+
inquirer.Path(
88+
"project_path",
89+
message=click.style("Enter the project directory path", fg="blue", bold=True),
90+
path_type=inquirer.Path.DIRECTORY,
91+
exists=False,
92+
),
93+
inquirer.Confirm(
94+
"custom_network",
95+
message=click.style(
96+
"Do you want to create a custom network?", fg="blue", bold=True
97+
),
98+
default=True,
99+
),
100+
]
101+
proj_answers = inquirer.prompt(questions)
102+
if proj_answers is None:
103+
click.secho("Setup cancelled by user.", fg="yellow")
104+
return False
105+
if not proj_answers["custom_network"]:
106+
project_path = Path(os.path.expanduser(proj_answers["project_path"]))
107+
create_warnet_project(project_path)
89108
click.secho("\nSetup completed successfully!", fg="green", bold=True)
109+
click.echo(
110+
"\nRun the following command to deploy this network using the default demo network:"
111+
)
112+
click.echo(f"warcli deploy {proj_answers['project_path']}/networks/6_node_bitcoin")
90113
return True
91-
92-
network_name = click.prompt(
93-
click.style("\nEnter the network name", fg="blue", bold=True),
94-
type=str,
95-
)
96-
97-
nodes = click.prompt(
98-
click.style("\nHow many nodes would you like?", fg="blue", bold=True),
99-
type=int,
100-
default=15,
101-
)
102-
connections = click.prompt(
103-
click.style(
104-
"\nHow many connections would you like each node to have?", fg="blue", bold=True
114+
answers.update(proj_answers)
115+
116+
# Custom network configuration
117+
questions = [
118+
inquirer.Text(
119+
"network_name",
120+
message=click.style("Enter your network name", fg="blue", bold=True),
121+
validate=lambda _, x: len(x) > 0,
105122
),
106-
type=int,
107-
default=8,
108-
)
109-
version = click.prompt(
110-
click.style(
111-
"\nWhich version would you like nodes to be by default?", fg="blue", bold=True
123+
inquirer.List(
124+
"nodes",
125+
message=click.style("How many nodes would you like?", fg="blue", bold=True),
126+
choices=["8", "12", "20", "50", "other"],
127+
default="12",
112128
),
113-
type=click.Choice(SUPPORTED_TAGS, case_sensitive=False),
114-
default=DEFAULT_TAG,
115-
)
129+
inquirer.List(
130+
"connections",
131+
message=click.style(
132+
"How many connections would you like each node to have?",
133+
fg="blue",
134+
bold=True,
135+
),
136+
choices=["0", "1", "2", "8", "12", "other"],
137+
default="8",
138+
),
139+
inquirer.List(
140+
"version",
141+
message=click.style(
142+
"Which version would you like nodes to run by default?", fg="blue", bold=True
143+
),
144+
choices=SUPPORTED_TAGS,
145+
default=DEFAULT_TAG,
146+
),
147+
]
148+
149+
net_answers = inquirer.prompt(questions)
150+
if net_answers is None:
151+
click.secho("Setup cancelled by user.", fg="yellow")
152+
return False
153+
154+
if net_answers["nodes"] == "other":
155+
custom_nodes = inquirer.prompt(
156+
[
157+
inquirer.Text(
158+
"nodes",
159+
message=click.style("Enter the number of nodes", fg="blue", bold=True),
160+
validate=lambda _, x: int(x) > 0,
161+
)
162+
]
163+
)
164+
if custom_nodes is None:
165+
click.secho("Setup cancelled by user.", fg="yellow")
166+
return False
167+
net_answers["nodes"] = custom_nodes["nodes"]
168+
169+
if net_answers["connections"] == "other":
170+
custom_connections = inquirer.prompt(
171+
[
172+
inquirer.Text(
173+
"connections",
174+
message=click.style(
175+
"Enter the number of connections", fg="blue", bold=True
176+
),
177+
validate=lambda _, x: int(x) >= 0,
178+
)
179+
]
180+
)
181+
if custom_connections is None:
182+
click.secho("Setup cancelled by user.", fg="yellow")
183+
return False
184+
net_answers["connections"] = custom_connections["connections"]
185+
answers.update(net_answers)
116186

117187
click.secho("\nCreating project structure...", fg="yellow", bold=True)
118-
create_warnet_project(Path(project_path))
188+
project_path = Path(os.path.expanduser(proj_answers["project_path"]))
189+
create_warnet_project(project_path)
119190
click.secho("\nGenerating custom network...", fg="yellow", bold=True)
120-
custom_network_path = Path(project_path) / "networks" / network_name
121-
custom_graph(nodes, connections, version, custom_network_path)
191+
custom_network_path = project_path / "networks" / answers["network_name"]
192+
custom_graph(
193+
int(answers["nodes"]),
194+
int(answers["connections"]),
195+
answers["version"],
196+
custom_network_path,
197+
)
122198
click.secho("\nSetup completed successfully!", fg="green", bold=True)
123199
click.echo("\nRun the following command to deploy this network:")
124200
click.echo(f"warnet deploy {custom_network_path}")
125201
except Exception as e:
202+
click.echo(f"{e}\n\n")
126203
click.secho(f"An error occurred while running the quick start script:\n\n{e}\n\n", fg="red")
127204
click.secho(
128-
"Please report this to https://github.com/bitcoin-dev-project/warnet/issues",
205+
"Please report the above context to https://github.com/bitcoin-dev-project/warnet/issues",
129206
fg="yellow",
130207
)
131208
return False
@@ -286,13 +363,15 @@ def custom_graph(num_nodes: int, num_connections: int, version: str, datadir: Pa
286363
datadir.mkdir(parents=False, exist_ok=False)
287364
# Generate network.yaml
288365
nodes = []
366+
connections = set()
289367

290368
for i in range(num_nodes):
291369
node = {"name": f"tank-{i:04d}", "connect": [], "image": {"tag": version}}
292370

293371
# Add round-robin connection
294372
next_node = (i + 1) % num_nodes
295373
node["connect"].append(f"tank-{next_node:04d}")
374+
connections.add((i, next_node))
296375

297376
# Add random connections
298377
available_nodes = list(range(num_nodes))
@@ -302,8 +381,11 @@ def custom_graph(num_nodes: int, num_connections: int, version: str, datadir: Pa
302381

303382
for _ in range(min(num_connections - 1, len(available_nodes))):
304383
random_node = random.choice(available_nodes)
305-
node["connect"].append(f"tank-{random_node:04d}")
306-
available_nodes.remove(random_node)
384+
# Avoid circular loops of A -> B -> A
385+
if (random_node, i) not in connections:
386+
node["connect"].append(f"tank-{random_node:04d}")
387+
connections.add((i, random_node))
388+
available_nodes.remove(random_node)
307389

308390
nodes.append(node)
309391

0 commit comments

Comments
 (0)