Utils - Snailz Docs

Skip to content

# Utils

Data generation utilities.

## `ensure_id_generator(cls)`

Ensure class has ID generator.

Source code in `src/snailz/utils.py`

    12
    13
    14
    15
    16
    def ensure_id_generator(cls):
        """Ensure class has ID generator."""

        if not hasattr(cls, "_id_gen"):
            cls._id_gen = id_gen(cls.id_stem, cls.id_digits)

## `file_or_std(parent, filename, mode)`

Open file and return handle or return stdin/stdout.

Source code in `src/snailz/utils.py`

    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    @contextmanager
    def file_or_std(parent, filename, mode):
        """Open file and return handle or return stdin/stdout."""

        if parent:
            stream = open(Path(parent, filename), mode)
            try:
                yield stream
            finally:
                stream.close()
        elif mode == "r":
            yield sys.stdin
        elif mode == "w":
            yield sys.stdout
        else:
            raise ValueError(f"bad filename/mode '{filename}' / '{mode}'")

## `id_gen(stem, digits)`

Generate unique IDs of the form 'stemDDDD'.

Source code in `src/snailz/utils.py`

    37
    38
    39
    40
    41
    42
    43
    44
    45
    def id_gen(stem, digits):
        """Generate unique IDs of the form 'stemDDDD'."""

        i = 1
        while True:
            temp = str(i)
            assert len(temp) <= digits, f"ID generation overflow {stem}: {i}"
            yield f"{stem}{temp.zfill(digits)}"
            i += 1

## `json_dump(obj, indent=2)`

Dump as JSON with custom serializer.

Source code in `src/snailz/utils.py`

    48
    49
    50
    51
    def json_dump(obj, indent=2):
        """Dump as JSON with custom serializer."""

        return json.dumps(obj, indent=indent, default=_serialize_json)

## `random_date(params)`

Select random date in range (inclusive).

Source code in `src/snailz/utils.py`

    54
    55
    56
    57
    58
    def random_date(params):
        """Select random date in range (inclusive)."""

        days = (params.sample_date[1] - params.sample_date[0]).days
        return params.sample_date[0] + timedelta(days=random.randint(0, days))

## `random_mass(params)`

Generate random sample mass.

Source code in `src/snailz/utils.py`

    61
    62
    63
    64
    65
    66
    67
    def random_mass(params):
        """Generate random sample mass."""

        return random.uniform(
            params.sample_mass[0],
            params.sample_mass[1],
        )

## `_serialize_json(obj)`

Custom JSON serializer.

Source code in `src/snailz/utils.py`

    70
    71
    72
    73
    74
    75
    76
    77
    def _serialize_json(obj):
        """Custom JSON serializer."""

        if isinstance(obj, date):
            return obj.isoformat()
        if isinstance(obj, BaseModel):
            return obj.model_dump()
        raise TypeError(f"Type {type(obj)} not serializable")