communex.cli._common

  1from dataclasses import dataclass
  2from getpass import getpass
  3from typing import Any, Mapping, cast, TypeVar
  4
  5import rich
  6import typer
  7from rich import box
  8from rich.console import Console
  9from rich.table import Table
 10from typer import Context
 11
 12from communex._common import get_node_url
 13from communex.balance import from_horus, from_nano, dict_from_nano
 14from communex.client import CommuneClient
 15from communex.types import (
 16    ModuleInfoWithOptionalBalance,
 17    NetworkParams,
 18    SubnetParamsWithEmission,
 19    )
 20
 21
 22@dataclass
 23class ExtraCtxData:
 24    output_json: bool
 25    use_testnet: bool
 26    yes_to_all: bool
 27
 28
 29class ExtendedContext(Context):
 30    obj: ExtraCtxData
 31
 32
 33@dataclass
 34class CustomCtx:
 35    ctx: ExtendedContext
 36    console: rich.console.Console
 37    console_err: rich.console.Console
 38    _com_client: CommuneClient | None = None
 39
 40    def com_client(self) -> CommuneClient:
 41        use_testnet = self.ctx.obj.use_testnet
 42        if self._com_client is None:
 43            node_url = get_node_url(None, use_testnet=use_testnet)
 44            self.info(f"Using node: {node_url}")
 45            for _ in range(5):
 46                try:
 47                    self._com_client = CommuneClient(
 48                        url=node_url, num_connections=1, wait_for_finalization=False)
 49                except Exception:
 50                    self.info(f"Failed to connect to node: {node_url}")
 51                    node_url = get_node_url(None, use_testnet=use_testnet)
 52                    self.info(f"Will retry with node {node_url}")
 53                    continue
 54            if self._com_client is None:
 55                raise ConnectionError("Could not connect to any node")
 56
 57        return self._com_client
 58
 59    def get_use_testnet(self) -> bool:
 60        return self.ctx.obj.use_testnet
 61
 62    def output(
 63        self,
 64        message: str,
 65        *args: tuple[Any, ...],
 66        **kwargs: dict[str, Any],
 67    ) -> None:
 68        self.console.print(message, *args, **kwargs)  # type: ignore
 69
 70    def info(
 71        self,
 72        message: str,
 73        *args: tuple[Any, ...],
 74        **kwargs: dict[str, Any],
 75    ) -> None:
 76
 77        self.console_err.print(message, *args, **kwargs)  # type: ignore
 78
 79    def error(self, message: str) -> None:
 80        message = f"ERROR: {message}"
 81        self.console_err.print(message, style="bold red")
 82
 83    def progress_status(self, message: str):
 84        return self.console_err.status(message)
 85
 86    def confirm(self, message: str) -> bool:
 87        if (self.ctx.obj.yes_to_all):
 88            print(f"{message} (--yes)")
 89            return True
 90        return typer.confirm(message)
 91
 92
 93def make_custom_context(ctx: typer.Context) -> CustomCtx:
 94    return CustomCtx(
 95        ctx=cast(ExtendedContext, ctx),
 96        console=Console(),
 97        console_err=Console(stderr=True),
 98    )
 99
100
101# Formatting
102
103
104def eprint(e: Any) -> None:
105    """
106    Pretty prints an error.
107    """
108
109    console = Console()
110
111    console.print(f"[bold red]ERROR: {e}", style="italic")
112
113
114def print_table_from_plain_dict(
115    result: Mapping[str, str | int | float | dict[Any, Any]], column_names: list[str], console: Console
116) -> None:
117    """
118    Creates a table for a plain dictionary.
119    """
120
121    table = Table(show_header=True, header_style="bold magenta")
122
123    for name in column_names:
124        table.add_column(name, style="white", vertical="middle")
125
126    # Add non-dictionary values to the table first
127    for key, value in result.items():
128        if not isinstance(value, dict):
129            table.add_row(key, str(value))
130    # Add subtables for nested dictionaries.
131    # Important to add after so that the display of the table is nicer.
132    for key, value in result.items():
133        if isinstance(value, dict):
134            subtable = Table(show_header=False, padding=(0, 0, 0, 0), border_style="bright_black")
135            for subkey, subvalue in value.items():
136                subtable.add_row(f"{subkey}: {subvalue}")
137            table.add_row(key, subtable)
138
139    console.print(table)
140
141
142
143def print_table_standardize(result: dict[str, list[Any]], console: Console) -> None:
144    """
145    Creates a table for a standardized dictionary.
146    """
147    table = Table(show_header=True, header_style="bold magenta")
148
149    for key in result.keys():
150        table.add_column(key, style="white")
151    rows = [*result.values()]
152    zipped_rows = [list(column) for column in zip(*rows)]
153    for row in zipped_rows:
154        table.add_row(*row, style="white")
155
156    console.print(table)
157
158
159def transform_module_into(
160    to_exclude: list[str], last_block: int,
161    immunity_period: int, modules: list[ModuleInfoWithOptionalBalance],
162    tempo: int
163):
164    mods = cast(list[dict[str, Any]], modules)
165    transformed_modules: list[dict[str, Any]] = []
166    for mod in mods:
167        module = mod.copy()
168        module_regblock = module["regblock"]
169        module["in_immunity"] = module_regblock + immunity_period > last_block
170
171        for key in to_exclude:
172            del module[key]
173        module["stake"] = round(from_nano(module["stake"]), 2)  # type: ignore
174        module["emission"] = round(
175            from_horus(
176                module["emission"], tempo
177            ),
178            4
179        )  # type: ignore
180        if module.get("balance") is not None:
181            module["balance"] = from_nano(module["balance"])  # type: ignore
182        else:
183            # user should not see None values
184            del module["balance"]
185        transformed_modules.append(module)
186
187    return transformed_modules
188
189
190def print_module_info(
191        client: CommuneClient,
192        modules: list[ModuleInfoWithOptionalBalance],
193        console: Console,
194        netuid: int,
195        title: str | None = None,
196) -> None:
197    """
198    Prints information about a module.
199    """
200    if not modules:
201        return
202
203    # Get the current block number, we will need this to caluclate immunity period
204    block = client.get_block()
205    if block:
206        last_block = block["header"]["number"]
207    else:
208        raise ValueError("Could not get block info")
209
210    # Get the immunity period on the netuid
211    immunity_period = client.get_immunity_period(netuid)
212    tempo = client.get_tempo(netuid)
213
214    # Transform the module dictionary to have immunity_period
215    table = Table(
216        show_header=True, header_style="bold magenta",
217        box=box.DOUBLE_EDGE, title=title,
218        caption_style="chartreuse3",
219        title_style="bold magenta",
220
221    )
222
223    to_exclude = ["stake_from", "last_update", "regblock"]
224    tranformed_modules = transform_module_into(
225        to_exclude, last_block, immunity_period, modules, tempo
226    )
227
228    sample_mod = tranformed_modules[0]
229    for key in sample_mod.keys():
230        # add columns
231        table.add_column(key, style="white")
232
233    total_stake = 0
234    total_balance = 0
235
236    for mod in tranformed_modules:
237        total_stake += mod["stake"]
238        if mod.get("balance") is not None:
239            total_balance += mod["balance"]
240
241        row: list[str] = []
242        for val in mod.values():
243            row.append(str(val))
244        table.add_row(*row)
245
246    table.caption = "total balance: " + f"{total_balance + total_stake}J"
247    console.print(table)
248    for _ in range(3):
249        console.print()
250
251
252def get_universal_password(ctx: CustomCtx) -> str:
253    ctx.info("Please provide the universal password for all keys")
254    universal_password = getpass()
255    return universal_password
256
257
258def tranform_network_params(params: NetworkParams):
259    """Transform network params to be human readable."""
260    governance_config = params["governance_config"]
261    allocation = governance_config["proposal_reward_treasury_allocation"]
262    governance_config = cast(dict[str, Any], governance_config)
263    governance_config["proposal_reward_treasury_allocation"] = f"{allocation}%"
264    params_ = cast(dict[str, Any], params)
265    params_["governance_config"] = governance_config
266    general_params = dict_from_nano(params_, [
267        "min_weight_stake",
268        "general_subnet_application_cost",
269        "subnet_registration_cost",
270        "proposal_cost",
271        "max_proposal_reward_treasury_allocation",
272    ])
273
274    return general_params
275
276
277T = TypeVar("T")
278V = TypeVar("V")
279def remove_none_values(data: dict[T, V | None]) -> dict[T, V]:
280    """
281    Removes key-value pairs from a dictionary where the value is None.
282    Works recursively for nested dictionaries.
283    """
284    cleaned_data: dict[T, V] = {}
285    for key, value in data.items():
286        if isinstance(value, dict):
287            cleaned_value = remove_none_values(value) # type: ignore
288            if cleaned_value is not None: # type: ignore
289                cleaned_data[key] = cleaned_value
290        elif value is not None:
291            cleaned_data[key] = value
292    return cleaned_data
293
294
295
296def transform_subnet_params(params: SubnetParamsWithEmission):
297    """Transform subnet params to be human readable."""
298    params_ = cast(dict[str, Any], params)
299    display_params = remove_none_values(params_)
300    display_params = dict_from_nano(
301        display_params, [
302            "bonds_ma",
303            "min_burn",
304            "max_burn",
305            "min_weight_stake",
306            "proposal_cost",
307            "max_proposal_reward_treasury_allocation",
308        ]
309    )
310    return display_params