Skip to content

Utils

Data generation utilities.

file_or_std(parent, filename, mode)

Open file and return handle or return stdin/stdout.

Source code in src/snailz/utils.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
@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 
30
31
32
33
34
35
36
37
38
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 
41
42
43
44
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 
47
48
49
50
51
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_size(params)

Generate random sample mass and diameter.

Source code in src/snailz/utils.py 
54
55
56
57
58
59
def random_size(params):
    """Generate random sample mass and diameter."""

    mass = random.normalvariate(*params.sample_mass)
    diameter = random.normalvariate(mass / 2.0, params.sample_mass[1] / 5.0)
    return mass, diameter

_serialize_json(obj)

Custom JSON serializer.

Source code in src/snailz/utils.py 
62
63
64
65
66
67
68
69
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")