Files
chia-blockchain/benchmarks/utils.py
T
6050235bdb [CHIA-3043] Add (hopefully) non-controversial ruff rules (#19684)
* Add (hopefully) non-controversial ruff rules

* mypy-exclusions.txt

* PIE

* Fix pre-commit?

* PYI

* RSE

* S + ignores

* Whitespace fix

* Re-add build-init-files.py

* Add comment about security ignores

* Update mypy-exclusions.txt

Co-authored-by: Kyle Altendorf <sda@fstab.net>

* Fix the executable status of build-init-files.py

* Set executable bit for tools/run_block.py and readd shebang

* Fix new errors

---------

Co-authored-by: Kyle Altendorf <sda@fstab.net>
2025-06-11 08:36:08 -07:00

71 lines
2.1 KiB
Python

from __future__ import annotations
import contextlib
import enum
import os
import subprocess
import sys
from collections.abc import AsyncIterator
from pathlib import Path
from typing import Any, Generic, Optional, TypeVar, Union
import click
from chia.util.db_wrapper import DBWrapper2
_T_Enum = TypeVar("_T_Enum", bound=enum.Enum)
# Workaround to allow `Enum` with click.Choice: https://github.com/pallets/click/issues/605#issuecomment-901099036
class EnumType(click.Choice, Generic[_T_Enum]):
def __init__(self, enum: type[_T_Enum], case_sensitive: bool = False) -> None:
self.__enum = enum
super().__init__(choices=[item.value for item in enum], case_sensitive=case_sensitive)
def convert(self, value: Any, param: Optional[click.Parameter], ctx: Optional[click.Context]) -> _T_Enum:
converted_str = super().convert(value, param, ctx)
return self.__enum(converted_str)
@contextlib.asynccontextmanager
async def setup_db(name: Union[str, os.PathLike[str]], db_version: int) -> AsyncIterator[DBWrapper2]:
db_filename = Path(name)
try:
os.unlink(db_filename)
except FileNotFoundError:
pass
log_path: Optional[Path]
if "--sql-logging" in sys.argv:
log_path = Path("sql.log")
else:
log_path = None
async with DBWrapper2.managed(
database=db_filename,
log_path=log_path,
db_version=db_version,
reader_count=1,
journal_mode="wal",
synchronous="full",
) as db_wrapper:
yield db_wrapper
def get_commit_hash() -> str:
try:
os.chdir(Path(os.path.realpath(__file__)).parent)
commit_hash = (
subprocess.run(["git", "rev-parse", "--short", "HEAD"], check=True, stdout=subprocess.PIPE)
.stdout.decode("utf-8")
.strip()
)
except Exception:
sys.exit("Failed to get the commit hash")
try:
if len(subprocess.run(["git", "status", "-s"], check=True, stdout=subprocess.PIPE).stdout) > 0:
raise Exception
except Exception:
commit_hash += "-dirty"
return commit_hash