communex.cli._common

  1from dataclasses import dataclass
  2from getpass import getpass
  3from typing import Any, Callable, Mapping, TypeVar, cast
  4
  5import rich
  6import rich.prompt
  7import typer
  8from rich import box
  9from rich.console import Console
 10from rich.table import Table
 11from substrateinterface import Keypair
 12from typer import Context
 13
 14from communex._common import ComxSettings, get_node_url
 15from communex.balance import dict_from_nano, from_horus, from_nano
 16from communex.client import CommuneClient
 17from communex.compat.key import resolve_key_ss58_encrypted, try_classic_load_key
 18from communex.errors import InvalidPasswordError, PasswordNotProvidedError
 19from communex.types import (
 20    ModuleInfoWithOptionalBalance,
 21    NetworkParams,
 22    Ss58Address,
 23    SubnetParamsWithEmission,
 24)
 25
 26
 27@dataclass
 28class ExtraCtxData:
 29    output_json: bool
 30    use_testnet: bool
 31    yes_to_all: bool
 32
 33
 34class ExtendedContext(Context):
 35    obj: ExtraCtxData
 36
 37
 38class CliPasswordProvider:
 39    def __init__(
 40        self, settings: ComxSettings, prompt_secret: Callable[[str], str]
 41    ):
 42        self.settings = settings
 43        self.prompt_secret = prompt_secret
 44
 45    def get_password(self, key_name: str) -> str | None:
 46        key_map = self.settings.KEY_PASSWORDS
 47        if key_map is not None:
 48            password = key_map.get(key_name)
 49            if password is not None:
 50                return password.get_secret_value()
 51        # fallback to universal password
 52        password = self.settings.UNIVERSAL_PASSWORD
 53        if password is not None:
 54            return password.get_secret_value()
 55        else:
 56            return None
 57
 58    def ask_password(self, key_name: str) -> str:
 59        password = self.prompt_secret(
 60            f"Please provide the password for the key '{key_name}'"
 61        )
 62        return password
 63
 64
 65class CustomCtx:
 66    ctx: ExtendedContext
 67    settings: ComxSettings
 68    console: rich.console.Console
 69    console_err: rich.console.Console
 70    password_manager: CliPasswordProvider
 71    _com_client: CommuneClient | None = None
 72
 73    def __init__(
 74        self,
 75        ctx: ExtendedContext,
 76        settings: ComxSettings,
 77        console: rich.console.Console,
 78        console_err: rich.console.Console,
 79        com_client: CommuneClient | None = None,
 80    ):
 81        self.ctx = ctx
 82        self.settings = settings
 83        self.console = console
 84        self.console_err = console_err
 85        self._com_client = com_client
 86        self.password_manager = CliPasswordProvider(
 87            self.settings, self.prompt_secret
 88        )
 89
 90    def get_use_testnet(self) -> bool:
 91        return self.ctx.obj.use_testnet
 92
 93    def get_node_url(self) -> str:
 94        use_testnet = self.get_use_testnet()
 95        return get_node_url(self.settings, use_testnet=use_testnet)
 96
 97    def com_client(self) -> CommuneClient:
 98        if self._com_client is None:
 99            node_url = self.get_node_url()
100            self.info(f"Using node: {node_url}")
101            for _ in range(5):
102                try:
103                    self._com_client = CommuneClient(
104                        url=node_url,
105                        num_connections=1,
106                        wait_for_finalization=False,
107                        timeout=65,
108                    )
109                except Exception:
110                    self.info(f"Failed to connect to node: {node_url}")
111                    node_url = self.get_node_url()
112                    self.info(f"Will retry with node {node_url}")
113                    continue
114                else:
115                    break
116            if self._com_client is None:
117                raise ConnectionError("Could not connect to any node")
118
119        return self._com_client
120
121    def output(
122        self,
123        message: str,
124        *args: tuple[Any, ...],
125        **kwargs: dict[str, Any],
126    ) -> None:
127        self.console.print(message, *args, **kwargs)  # type: ignore
128
129    def info(
130        self,
131        message: str,
132        *args: tuple[Any, ...],
133        **kwargs: dict[str, Any],
134    ) -> None:
135        self.console_err.print(message, *args, **kwargs)  # type: ignore
136
137    def error(
138        self,
139        message: str,
140        *args: tuple[Any, ...],
141        **kwargs: dict[str, Any],
142    ) -> None:
143        message = f"ERROR: {message}"
144        self.console_err.print(message, *args, style="bold red", **kwargs)  # type: ignore
145
146    def progress_status(self, message: str):
147        return self.console_err.status(message)
148
149    def confirm(self, message: str) -> bool:
150        if self.ctx.obj.yes_to_all:
151            print(f"{message} (--yes)")
152            return True
153        return typer.confirm(message, err=True)
154
155    def prompt_secret(self, message: str) -> str:
156        return rich.prompt.Prompt.ask(
157            message, password=True, console=self.console_err
158        )
159
160    def load_key(self, key: str, password: str | None = None) -> Keypair:
161        try:
162            keypair = try_classic_load_key(
163                key, password, password_provider=self.password_manager
164            )
165            return keypair
166        except PasswordNotProvidedError:
167            self.error(f"Password not provided for key '{key}'")
168            raise typer.Exit(code=1)
169        except InvalidPasswordError:
170            self.error(f"Incorrect password for key '{key}'")
171            raise typer.Exit(code=1)
172
173    def resolve_key_ss58(
174        self, key: Ss58Address | Keypair | str, password: str | None = None
175    ) -> Ss58Address:
176        try:
177            address = resolve_key_ss58_encrypted(
178                key, password, password_provider=self.password_manager
179            )
180            return address
181        except PasswordNotProvidedError:
182            self.error(f"Password not provided for key '{key}'")
183            raise typer.Exit(code=1)
184        except InvalidPasswordError:
185            self.error(f"Incorrect password for key '{key}'")
186            raise typer.Exit(code=1)
187
188
189def make_custom_context(ctx: typer.Context) -> CustomCtx:
190    return CustomCtx(
191        ctx=cast(ExtendedContext, ctx),  # TODO: better check
192        settings=ComxSettings(),
193        console=Console(),
194        console_err=Console(stderr=True),
195    )
196
197
198# Formatting
199
200
201def eprint(e: Any) -> None:
202    """
203    Pretty prints an error.
204    """
205
206    console = Console()
207
208    console.print(f"[bold red]ERROR: {e}", style="italic")
209
210
211def print_table_from_plain_dict(
212    result: Mapping[str, str | int | float | dict[Any, Any] | Ss58Address],
213    column_names: list[str],
214    console: Console,
215) -> None:
216    """
217    Creates a table for a plain dictionary.
218    """
219
220    table = Table(show_header=True, header_style="bold magenta")
221
222    for name in column_names:
223        table.add_column(name, style="white", vertical="middle")
224
225    # Add non-dictionary values to the table first
226    for key, value in result.items():
227        if not isinstance(value, dict):
228            table.add_row(key, str(value))
229    # Add subtables for nested dictionaries.
230    # Important to add after so that the display of the table is nicer.
231    for key, value in result.items():
232        if isinstance(value, dict):
233            subtable = Table(
234                show_header=False,
235                padding=(0, 0, 0, 0),
236                border_style="bright_black",
237            )
238            for subkey, subvalue in value.items():
239                subtable.add_row(f"{subkey}: {subvalue}")
240            table.add_row(key, subtable)
241
242    console.print(table)
243
244
245def print_table_standardize(
246    result: dict[str, list[Any]], console: Console
247) -> None:
248    """
249    Creates a table for a standardized dictionary.
250    """
251    table = Table(show_header=True, header_style="bold magenta")
252
253    for key in result.keys():
254        table.add_column(key, style="white")
255    rows = [*result.values()]
256    zipped_rows = [list(column) for column in zip(*rows)]
257    for row in zipped_rows:
258        table.add_row(*row, style="white")
259
260    console.print(table)
261
262
263def transform_module_into(
264    to_exclude: list[str],
265    last_block: int,
266    immunity_period: int,
267    modules: list[ModuleInfoWithOptionalBalance],
268    tempo: int,
269):
270    mods = cast(list[dict[str, Any]], modules)
271    transformed_modules: list[dict[str, Any]] = []
272    for mod in mods:
273        module = mod.copy()
274        module_regblock = module["regblock"]
275        module["in_immunity"] = module_regblock + immunity_period > last_block
276
277        for key in to_exclude:
278            del module[key]
279        module["stake"] = round(from_nano(module["stake"]), 2)  # type: ignore
280        module["emission"] = round(from_horus(module["emission"], tempo), 4)  # type: ignore
281        if module.get("balance") is not None:
282            module["balance"] = from_nano(module["balance"])  # type: ignore
283        else:
284            # user should not see None values
285            del module["balance"]
286        transformed_modules.append(module)
287
288    return transformed_modules
289
290
291def print_module_info(
292    client: CommuneClient,
293    modules: list[ModuleInfoWithOptionalBalance],
294    console: Console,
295    netuid: int,
296    title: str | None = None,
297) -> None:
298    """
299    Prints information about a module.
300    """
301    if not modules:
302        return
303
304    # Get the current block number, we will need this to caluclate immunity period
305    block = client.get_block()
306    if block:
307        last_block = block["header"]["number"]
308    else:
309        raise ValueError("Could not get block info")
310
311    # Get the immunity period on the netuid
312    immunity_period = client.get_immunity_period(netuid)
313    tempo = client.get_tempo(netuid)
314
315    # Transform the module dictionary to have immunity_period
316    table = Table(
317        show_header=True,
318        header_style="bold magenta",
319        box=box.DOUBLE_EDGE,
320        title=title,
321        caption_style="chartreuse3",
322        title_style="bold magenta",
323    )
324
325    to_exclude = ["stake_from", "last_update", "regblock"]
326    tranformed_modules = transform_module_into(
327        to_exclude, last_block, immunity_period, modules, tempo
328    )
329
330    sample_mod = tranformed_modules[0]
331    for key in sample_mod.keys():
332        # add columns
333        table.add_column(key, style="white")
334
335    total_stake = 0
336    total_balance = 0
337
338    for mod in tranformed_modules:
339        total_stake += mod["stake"]
340        if mod.get("balance") is not None:
341            total_balance += mod["balance"]
342
343        row: list[str] = []
344        for val in mod.values():
345            row.append(str(val))
346        table.add_row(*row)
347
348    table.caption = "total balance: " + f"{total_balance + total_stake}J"
349    console.print(table)
350    for _ in range(3):
351        console.print()
352
353
354def get_universal_password(ctx: CustomCtx) -> str:
355    ctx.info("Please provide the universal password for all keys")
356    universal_password = getpass()
357    return universal_password
358
359
360def tranform_network_params(params: NetworkParams):
361    """Transform network params to be human readable."""
362    governance_config = params["governance_config"]
363    allocation = governance_config["proposal_reward_treasury_allocation"]
364    governance_config = cast(dict[str, Any], governance_config)
365    governance_config["proposal_reward_treasury_allocation"] = f"{allocation}%"
366    params_ = cast(dict[str, Any], params)
367    params_["governance_config"] = governance_config
368    general_params = dict_from_nano(
369        params_,
370        [
371            "min_weight_stake",
372            "general_subnet_application_cost",
373            "subnet_registration_cost",
374            "proposal_cost",
375            "max_proposal_reward_treasury_allocation",
376        ],
377    )
378
379    return general_params
380
381
382T = TypeVar("T")
383V = TypeVar("V")
384
385
386def remove_none_values(data: dict[T, V | None]) -> dict[T, V]:
387    """
388    Removes key-value pairs from a dictionary where the value is None.
389    Works recursively for nested dictionaries.
390    """
391    cleaned_data: dict[T, V] = {}
392    for key, value in data.items():
393        if isinstance(value, dict):
394            cleaned_value = remove_none_values(value)  # type: ignore
395            if cleaned_value is not None:  # type: ignore
396                cleaned_data[key] = cleaned_value
397        elif value is not None:
398            cleaned_data[key] = value
399    return cleaned_data
400
401
402def transform_subnet_params(params: dict[int, SubnetParamsWithEmission]):
403    """Transform subnet params to be human readable."""
404    params_ = cast(dict[int, Any], params)
405    display_params = remove_none_values(params_)
406    display_params = dict_from_nano(
407        display_params,
408        [
409            "bonds_ma",
410            "min_burn",
411            "max_burn",
412            "min_weight_stake",
413            "proposal_cost",
414            "max_proposal_reward_treasury_allocation",
415            "min_validator_stake",
416        ],
417    )
418    return display_params