Skip to content

Utils

Data generation utilities.

ensure_id_generator(cls)

Ensure class has ID generator.

Source code in snailz/utils.py 
13
14
15
16
17
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 snailz/utils.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
@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 snailz/utils.py 
38
39
40
41
42
43
44
45
46
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 snailz/utils.py 
49
50
51
52
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 snailz/utils.py 
55
56
57
58
59
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 snailz/utils.py 
62
63
64
65
66
67
68
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 snailz/utils.py 
71
72
73
74
75
76
77
78
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")