Coverage for src/crimepy/utilities.py: 96%
71 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-30 18:01 +0000
1"""Small tools that the rest of crimepy shares.
3Contents:
4 data_folder: returns the folder where crimepy saves data (by default, a
5 folder of the data folder of amos's workspace).
6 to_bool: makes a boolean from a boolean or text.
7 to_dataset: makes an `amos.Dataset` from a dataset or a data frame.
8 to_frame: returns the data frame of a dataset or a data frame.
9 to_list: makes a list from one item or several.
10 to_names: makes a list of names from names or text separated by commas.
11 to_years: makes a list of years from years, ranges, and text.
13"""
15from __future__ import annotations
17import pathlib
18import re
19from collections.abc import Iterable
20from typing import Any
22import amos
23import pandas as pd
25from . import options
27# A range of years written as text, such as "2019-2023" or "2019:2023".
28_RANGE: re.Pattern[str] = re.compile(r'^\s*(\d{4})\s*[-:]\s*(\d{4})\s*$')
29# Text that means true or false.
30_TRUE: frozenset[str] = frozenset({'true', 'yes', 'y', '1', 'on'})
31_FALSE: frozenset[str] = frozenset({'false', 'no', 'n', '0', 'off', 'none', ''})
34""" Public Functions """
37def data_folder(folder: pathlib.Path | str | None = None) -> pathlib.Path:
38 """Returns the folder where crimepy saves data, making it if needed.
40 By default, it is a folder (`options.DATA_SUBFOLDER`, "crime") of the
41 data folder of amos's workspace, the folder that holds the "code" folder
42 (see `amos.clerks.workspace`), such as "projects/data/crime" for code in
43 "projects/code". So crimepy keeps its data outside every project, where
44 an amos project's `Clerk` looks for data too.
46 Args:
47 folder: the folder to use. Defaults to `None`, in which case it is
48 `options.DATA_FOLDER`, if it is set (as it is by the
49 "CRIMEPY_DATA" environment variable), and otherwise the
50 `options.DATA_SUBFOLDER` of the data folder of amos's workspace.
52 Returns:
53 The folder.
55 """
56 chosen = folder or options.DATA_FOLDER
57 if chosen is None:
58 chosen = (
59 amos.clerks.workspace() / amos.options._DATA_FOLDER
60 / options.DATA_SUBFOLDER)
61 path = pathlib.Path(chosen).expanduser()
62 path.mkdir(parents = True, exist_ok = True)
63 return path
66def to_bool(item: Any) -> bool:
67 """Returns `item` as a boolean.
69 Args:
70 item: a boolean, a number, or text (such as "true", "yes", "1",
71 "false", "no", or "0", as a settings file may give it).
73 Raises:
74 ValueError: if text is not a boolean.
76 Returns:
77 The boolean.
79 """
80 if isinstance(item, str):
81 text = item.strip().lower()
82 if text in _TRUE:
83 return True
84 if text in _FALSE:
85 return False
86 message = f'{item!r} is not true or false'
87 raise ValueError(message)
88 return bool(item)
91def to_dataset(item: Any, **kwargs: Any) -> amos.Dataset:
92 """Returns `item` as an `amos.Dataset`.
94 Args:
95 item: an `amos.Dataset` (which is returned as it is) or anything that
96 `amos.Dataset.create` accepts, such as a `pandas.DataFrame`.
97 **kwargs: arguments for `amos.Dataset.create`, such as "label" and
98 "groups". A dataset keeps any of them that it already has.
100 Returns:
101 The dataset.
103 """
104 return amos.Dataset.create(item, **kwargs)
107def to_frame(item: Any) -> pd.DataFrame:
108 """Returns the data frame of `item`.
110 Args:
111 item: an `amos.Dataset` or a `pandas.DataFrame`.
113 Returns:
114 The data of the dataset, or the data frame itself.
116 """
117 return item.data if isinstance(item, amos.Dataset) else item
120def to_list(item: Any) -> list[Any]:
121 """Returns `item` as a list.
123 Args:
124 item: `None`, one item, or an iterable of items. Text is one item.
126 Returns:
127 An empty list for `None`, the items of an iterable, and otherwise a
128 list with `item` in it.
130 """
131 if item is None:
132 return []
133 if isinstance(item, (str, bytes)) or not isinstance(item, Iterable):
134 return [item]
135 return list(item)
138def to_names(item: Any) -> list[str]:
139 """Returns `item` as a list of names.
141 Args:
142 item: `None`, a name, names separated by commas (as a settings file
143 may give them, such as "DE, MD"), or an iterable of names.
145 Returns:
146 The names, without blanks around them.
148 """
149 names: list[str] = []
150 for part in to_list(item):
151 names.extend(n.strip() for n in str(part).split(',') if n.strip())
152 return names
155def to_years(years: Any) -> list[int]:
156 """Returns a sorted list of years without duplicates.
158 Args:
159 years: a year, a `range`, text with years separated by commas or
160 spaces, a range of years as text ("2019-2023", which includes
161 both ends), or an iterable of any of them.
163 Raises:
164 ValueError: if a year is not a whole number.
166 Returns:
167 The years.
169 """
170 found: set[int] = set()
171 for item in to_list(years):
172 if isinstance(item, str):
173 for part in re.split(r'[,\s]+', item.strip()):
174 if not part: 174 ↛ 175line 174 didn't jump to line 175 because the condition on line 174 was never true
175 continue
176 matched = _RANGE.match(part)
177 if matched:
178 first, last = (int(y) for y in matched.groups())
179 found.update(range(min(first, last), max(first, last) + 1))
180 else:
181 found.add(_to_year(part))
182 elif isinstance(item, range):
183 found.update(item)
184 elif isinstance(item, Iterable): 184 ↛ 185line 184 didn't jump to line 185 because the condition on line 184 was never true
185 found.update(to_years(list(item)))
186 else:
187 found.add(_to_year(item))
188 return sorted(found)
191""" Private Functions """
194def _to_year(item: Any) -> int:
195 """Returns `item` as a year.
197 Args:
198 item: a year, as a number or text.
200 Raises:
201 ValueError: if `item` is not a whole number.
203 Returns:
204 The year.
206 """
207 if isinstance(item, float) and item.is_integer():
208 return int(item)
209 text = str(item).strip()
210 if not text.isdigit():
211 message = f'{item!r} is not a year'
212 raise ValueError(message)
213 return int(text)