Coverage for src/crimepy/nibrs/reader.py: 92%
273 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-30 18:01 +0000
1"""Reads and parses NIBRS master files.
3A master file is read one record at a time from its zip archive, without
4extracting it (recent files are over 6 GB of text). The records are sorted
5into tables by their segment level and parsed in chunks: each chunk of a
6segment becomes a grid of characters, and every field is sliced from the grid
7at once, so even the largest files are parsed in a few minutes.
9Each table has a column for each field of its layout (see
10`crimepy.nibrs.layouts`), and two more:
12* `data_year`: the year of the master file.
13* `agency_ori`: the ORI of the batch header that the record follows. The
14 records of an agency follow its batch header, and their ORIs usually match
15 it, but the FBI notes that the last two characters may differ.
17Pre-2013 files split each batch header into three records ("B1", "B2", and
18"B3"), which are joined into one row of the "agencies" table, so that it has
19the same columns in every year.
21Contents:
22 open_master: opens the text of a master file for reading.
23 parse: parses records of one segment level into a table.
24 read: reads tables from a master file into `amos` datasets.
25 records: yields the records of a master file.
26 tables: yields the tables of a master file in chunks.
27 year_of: finds the year of a master file from its name.
29"""
31from __future__ import annotations
33import collections
34import contextlib
35import datetime as dt
36import io
37import pathlib
38import re
39import struct
40import zipfile
41from collections.abc import Callable, Iterator, Sequence
42from typing import IO, Any, TypeAlias
44import amos
45import numpy as np
46import pyarrow as pa
48from .. import options, utilities
49from . import codes, layouts, munge
51# The compression method of zip archives that `zipfile` cannot read. The CDE
52# compressed some master files (such as those for 2017 to 2023 that it
53# published before August 2026) with it.
54_DEFLATE64: int = 9
55# The size of the fixed part of a zip archive's local file header.
56_LOCAL_HEADER: int = 30
57# Bytes of compressed data read at a time.
58_BLOCK: int = 1 << 22
59# A year in the name of a master file, such as "nibrs-2023.zip".
60_YEAR: re.Pattern[str] = re.compile(r'(?<!\d)(19[89]\d|20\d\d)(?!\d)')
61# A two-digit year in the name of an older master file, such as
62# "NIBRS04.ENC" or "KCNIBR91.DAT".
63_SHORT_YEAR: re.Pattern[str] = re.compile(
64 r'NIBRS?_?(?:YR)?(\d\d)(?!\d)', re.IGNORECASE)
65# The `pyarrow` type of each kind of field.
66_TYPES: dict[str, pa.DataType] = {
67 'bool': pa.bool_(),
68 'code': pa.string(),
69 'date': pa.date32(),
70 'int': pa.int64(),
71 'quantity': pa.float64(),
72 'text': pa.string(),
73 'yes': pa.bool_()}
74# The segment levels of batch headers.
75_HEADERS: frozenset[bytes] = frozenset({b'BH', b'B1'})
76# The segment levels of the three batch header records of pre-2013 files.
77_SPLIT_HEADERS: tuple[str, ...] = ('B1', 'B2', 'B3')
78# A chunk of the tables of a master file, by name.
79_Chunk: TypeAlias = dict[str, pa.Table]
82""" Private Classes """
85class _Deflate64(io.RawIOBase):
86 """Reads a member of a zip archive that is compressed with Deflate64.
88 Args:
89 path: path of the zip archive.
90 info: the member's entry in the archive.
92 """
94 """ Initialization Methods """
96 def __init__(self, path: pathlib.Path, info: zipfile.ZipInfo) -> None:
97 """Opens the archive at the start of the member's compressed data.
99 Args:
100 path: path of the zip archive.
101 info: the member's entry in the archive.
103 """
104 import inflate64 # noqa: PLC0415
106 super().__init__()
107 self._file = open(path, 'rb') # noqa: SIM115
108 self._file.seek(info.header_offset)
109 header = self._file.read(_LOCAL_HEADER)
110 name, extra = struct.unpack('<HH', header[26:30])
111 self._file.seek(name + extra, io.SEEK_CUR)
112 self._left = info.compress_size
113 self._inflater = inflate64.Inflater()
114 self._buffer = b''
116 """ Required Methods """
118 def readable(self) -> bool:
119 """Returns whether the member can be read.
121 Returns:
122 `True`, since the member can be read.
124 """
125 return True
127 def readinto(self, buffer: Any) -> int:
128 """Reads decompressed data into `buffer`.
130 Args:
131 buffer: a writable buffer.
133 Returns:
134 The number of bytes read, which is 0 at the end of the member.
136 """
137 while not self._buffer and self._left > 0:
138 block = self._file.read(min(_BLOCK, self._left))
139 if not block: 139 ↛ 140line 139 didn't jump to line 140 because the condition on line 139 was never true
140 break
141 self._left -= len(block)
142 self._buffer = self._inflater.inflate(block)
143 size = min(len(buffer), len(self._buffer))
144 buffer[:size] = self._buffer[:size]
145 self._buffer = self._buffer[size:]
146 return size
148 """ Public Methods """
150 def close(self) -> None:
151 """Closes the archive."""
152 self._file.close()
153 super().close()
156""" Public Functions """
159@contextlib.contextmanager
160def open_master(path: pathlib.Path | str) -> Iterator[IO[bytes]]:
161 """Opens the text of a master file for reading, as bytes.
163 Args:
164 path: path of a master file: a zip archive (as the CDE publishes it,
165 compressed with Deflate or Deflate64) or the text file in it.
167 Raises:
168 FileNotFoundError: if there is no file at `path`.
169 ValueError: if a zip archive is empty.
171 Yields:
172 IO[bytes]: the text of the master file.
174 """
175 path = pathlib.Path(path)
176 if not path.is_file():
177 message = f'there is no master file at {path}'
178 raise FileNotFoundError(message)
179 if not zipfile.is_zipfile(path):
180 with open(path, 'rb') as stream:
181 yield stream
182 return
183 with zipfile.ZipFile(path) as archive:
184 members = [i for i in archive.infolist() if not i.is_dir()]
185 if not members: 185 ↛ 186line 185 didn't jump to line 186 because the condition on line 185 was never true
186 message = f'the zip archive at {path} is empty'
187 raise ValueError(message)
188 info = max(members, key = lambda i: i.file_size)
189 if info.compress_type == _DEFLATE64:
190 raw = _Deflate64(path, info)
191 with io.BufferedReader(raw, buffer_size = _BLOCK) as stream:
192 yield stream
193 else:
194 with archive.open(info) as stream:
195 yield stream
198def records(path: pathlib.Path | str) -> Iterator[bytes]:
199 """Yields the records of a master file, without their line endings.
201 Args:
202 path: path of a master file (see `open_master`).
204 Yields:
205 bytes: each record.
207 """
208 with open_master(path) as stream:
209 for line in stream:
210 record = line.rstrip(b'\r\n')
211 if record: 211 ↛ 209line 211 didn't jump to line 209 because the condition on line 211 was always true
212 yield record
215def year_of(path: pathlib.Path | str) -> int | None:
216 """Returns the year of a master file, found from its name.
218 Args:
219 path: path of a master file, such as "nibrs-2023.zip".
221 Returns:
222 The year, or `None` if the name does not have one.
224 """
225 path = pathlib.Path(path)
226 found = _YEAR.search(path.name)
227 if found:
228 return int(found.group(1))
229 names = [path.name]
230 if zipfile.is_zipfile(path): 230 ↛ 231line 230 didn't jump to line 231 because the condition on line 230 was never true
231 with zipfile.ZipFile(path) as archive:
232 names.extend(archive.namelist())
233 for name in names:
234 found = _YEAR.search(name)
235 if found: 235 ↛ 236line 235 didn't jump to line 236 because the condition on line 235 was never true
236 return int(found.group(1))
237 short = _SHORT_YEAR.search(name)
238 if short:
239 year = int(short.group(1))
240 return year + (1900 if year >= 90 else 2000)
241 return None
244def parse(
245 lines: Sequence[bytes],
246 layout: layouts.Layout | str,
247 **constants: Any) -> pa.Table:
248 """Parses records of one segment level into a table.
250 Args:
251 lines: the records, as bytes, all of the level of `layout`.
252 layout: the layout of the records, or its segment level (such as
253 "04").
254 **constants: columns to add before the fields, each with one value
255 for every record (such as `data_year = 2023`), or a sequence of
256 values, one for each record.
258 Returns:
259 A table with a column for each constant and each field.
261 """
262 if isinstance(layout, str): 262 ↛ 264line 262 didn't jump to line 264 because the condition on line 262 was always true
263 layout = layouts.LAYOUTS[layout]
264 columns: dict[str, pa.Array] = {}
265 size = len(lines)
266 for name, value in constants.items():
267 columns[name] = _constant(value, size)
268 if size == 0:
269 for field in layout:
270 columns[field.name] = pa.array([], type = _TYPES[field.kind])
271 return pa.table(columns)
272 width = layout.width
273 text = b''.join(line[:width].ljust(width) for line in lines)
274 grid = np.frombuffer(text, dtype = np.uint8).reshape(size, width)
275 for field in layout:
276 cells = np.ascontiguousarray(grid[:, field.start:field.end])
277 columns[field.name] = _convert(
278 cells.view(f'S{field.width}').ravel(), field)
279 return pa.table(columns)
282def schema(table: str) -> pa.Schema:
283 """Returns the columns and types of a table of a master file.
285 Args:
286 table: the name of the table (such as "victims").
288 Raises:
289 KeyError: if `table` is not a table of a master file.
291 Returns:
292 The schema of the table.
294 """
295 if table not in layouts.SEGMENTS: 295 ↛ 296line 295 didn't jump to line 296 because the condition on line 295 was never true
296 message = (
297 f'{table!r} is not a table: use one of '
298 f'{sorted(layouts.SEGMENTS)}')
299 raise KeyError(message)
300 if table == 'agencies':
301 fields = list(layouts.LAYOUTS['BH'])
302 extra = [('data_year', pa.int64())]
303 else:
304 fields = list(layouts.LAYOUTS[layouts.SEGMENTS[table][0]])
305 extra = [('data_year', pa.int64()), ('agency_ori', pa.string())]
306 return pa.schema(extra + [(f.name, _TYPES[f.kind]) for f in fields])
309def tables(
310 path: pathlib.Path | str,
311 names: Sequence[str] | str | None = None,
312 *,
313 year: int | None = None,
314 states: Sequence[str] | str | None = None,
315 chunk_size: int | None = None,
316 limit: int | None = None,
317 progress: Callable[[int], None] | None = None) -> Iterator[_Chunk]:
318 """Yields the tables of a master file, a chunk of records at a time.
320 Args:
321 path: path of a master file (see `open_master`).
322 names: the tables to make (see `crimepy.nibrs.layouts.TABLES`).
323 Defaults to `None`, which makes them all.
324 year: the year of the master file. Defaults to `None`, in which case
325 it is found from the file's name (see `year_of`).
326 states: postal abbreviations (such as "DE") or numeric state codes
327 (such as "07") of the states to keep. Defaults to `None`, which
328 keeps every state.
329 chunk_size: the number of records to parse at once. Defaults to
330 `None`, in which case it is `crimepy.options.CHUNK_SIZE`.
331 limit: the number of records to read before stopping. Defaults to
332 `None`, which reads them all.
333 progress: called with the number of records read so far after each
334 chunk. Defaults to `None`.
336 Yields:
337 dict[str, pa.Table]: for each chunk, the tables that have records in
338 it, by name. The "agencies" table, which is small, comes in the
339 last chunk.
341 """
342 wanted = _wanted(names)
343 levels = {
344 level.encode() for name in wanted
345 for level in layouts.SEGMENTS[name]}
346 kept = _state_codes(states)
347 year = year if year is not None else year_of(path)
348 chunk_size = chunk_size or options.CHUNK_SIZE
349 batch: dict[bytes, list[bytes]] = collections.defaultdict(list)
350 owners: dict[bytes, list[bytes]] = collections.defaultdict(list)
351 headers: dict[str, list[bytes]] = collections.defaultdict(list)
352 owner = b''
353 count = 0
354 pending = 0
355 for record in records(path):
356 if limit is not None and count >= limit:
357 break
358 count += 1
359 level = record[:2]
360 if level in _HEADERS:
361 owner = record[4:13]
362 if level not in levels or (kept and record[2:4] not in kept):
363 continue
364 if level[0:1] == b'B':
365 headers[level.decode()].append(record)
366 continue
367 batch[level].append(record)
368 owners[level].append(owner)
369 pending += 1
370 if pending >= chunk_size:
371 yield _parse_batch(batch, owners, year)
372 batch.clear()
373 owners.clear()
374 pending = 0
375 if progress is not None: 375 ↛ 355line 375 didn't jump to line 355 because the condition on line 375 was always true
376 progress(count)
377 parsed = _parse_batch(batch, owners, year)
378 if 'agencies' in wanted:
379 parsed['agencies'] = _agencies(headers, year)
380 if progress is not None:
381 progress(count)
382 yield parsed
385def read(
386 path: pathlib.Path | str,
387 names: Sequence[str] | str | None = None,
388 *,
389 year: int | None = None,
390 states: Sequence[str] | str | None = None,
391 decode: bool = True,
392 chunk_size: int | None = None,
393 limit: int | None = None) -> dict[str, amos.Dataset]:
394 """Reads tables from a master file into `amos` datasets.
396 Everything is kept in memory, so reading a whole recent year (tens of
397 millions of records) needs a lot of it. To work with whole years, convert
398 them to parquet files once with `crimepy.nibrs.convert` and load what you
399 need with `crimepy.nibrs.load`. `read` suits a few states, a few tables,
400 or a quick look (with `limit`).
402 Args:
403 path: path of a master file (see `open_master`).
404 names: the tables to read (see `crimepy.nibrs.layouts.TABLES`).
405 Defaults to `None`, which reads them all.
406 year: the year of the master file. Defaults to `None`, in which case
407 it is found from the file's name.
408 states: postal abbreviations or numeric state codes of the states to
409 keep. Defaults to `None`, which keeps every state.
410 decode: whether to replace codes with their labels (see
411 `crimepy.nibrs.munge.prepare`). Defaults to `True`.
412 chunk_size: the number of records to parse at once. Defaults to
413 `None`, in which case it is `crimepy.options.CHUNK_SIZE`.
414 limit: the number of records to read before stopping. Defaults to
415 `None`, which reads them all.
417 Returns:
418 A dataset for each table, by name, whose `data` is the table.
420 """
421 wanted = _wanted(names)
422 parts: dict[str, list[pa.Table]] = {name: [] for name in wanted}
423 for chunk in tables(
424 path, wanted, year = year, states = states, chunk_size = chunk_size,
425 limit = limit):
426 for name, table in chunk.items():
427 parts[name].append(table)
428 datasets = {}
429 for name, pieces in parts.items():
430 table = (
431 pa.concat_tables(pieces) if pieces
432 else schema(name).empty_table())
433 frame = munge.prepare(munge.to_pandas(table), name, decode = decode)
434 datasets[name] = amos.Dataset.create(frame)
435 return datasets
438""" Private Functions """
441def _agencies(headers: dict[str, list[bytes]], year: int | None) -> pa.Table:
442 """Returns the table of agencies from their batch header records.
444 Args:
445 headers: the batch header records, by segment level.
446 year: the year of the master file.
448 Returns:
449 The table of agencies, with the columns of `schema('agencies')`.
451 """
452 target = schema('agencies')
453 if headers.get('BH'):
454 table = parse(headers['BH'], 'BH', data_year = year)
455 elif headers.get('B1'): 455 ↛ 467line 455 didn't jump to line 467 because the condition on line 455 was always true
456 table = parse(headers['B1'], 'B1', data_year = year)
457 for level in _SPLIT_HEADERS[1:]:
458 other = parse(headers.get(level, []), level)
459 other = other.group_by('ori', use_threads = False).aggregate(
460 [(c, 'first') for c in other.column_names if c != 'ori'])
461 other = other.rename_columns(
462 [c.removesuffix('_first') for c in other.column_names])
463 table = table.join(other, 'ori', join_type = 'left outer')
464 table = table.sort_by(
465 [('state_code', 'ascending'), ('ori', 'ascending')])
466 else:
467 return target.empty_table()
468 columns = [
469 table.column(f.name).cast(f.type) if f.name in table.column_names
470 else pa.nulls(len(table), type = f.type)
471 for f in target]
472 return pa.Table.from_arrays(columns, schema = target)
475def _constant(value: Any, size: int) -> pa.Array:
476 """Returns a column with `value` for each of `size` rows.
478 Args:
479 value: one value for every row, or a sequence of one value for each
480 row (such as the ORIs of their batch headers, as bytes).
481 size: the number of rows.
483 Returns:
484 The column.
486 """
487 if isinstance(value, (list, tuple, np.ndarray)):
488 cells = np.asarray(value, dtype = 'S9') if value else np.array([], 'S9')
489 field = layouts.Field('constant', 0, 9, 'text')
490 return _convert(cells, field)
491 if value is None: 491 ↛ 492line 491 didn't jump to line 492 because the condition on line 491 was never true
492 return pa.nulls(size, type = pa.int64())
493 return pa.array(np.full(size, value))
496def _convert(cells: np.ndarray, field: layouts.Field) -> pa.Array:
497 """Returns a column of the values of `field` in `cells`.
499 Each distinct value is converted once, so fields with few values (most
500 of them) are converted quickly.
502 Args:
503 cells: the text of the field in each record, as fixed-width bytes.
504 field: the field.
506 Returns:
507 The column, with the type of the field's kind.
509 """
510 uniques, inverse = np.unique(cells, return_inverse = True)
511 converter = _CONVERTERS[field.kind]
512 values = [converter(u, field) for u in uniques.tolist()]
513 array = pa.array(values, type = _TYPES[field.kind])
514 return array.take(pa.array(inverse.ravel().astype(np.int32)))
517def _parse_batch(
518 batch: dict[bytes, list[bytes]],
519 owners: dict[bytes, list[bytes]],
520 year: int | None) -> dict[str, pa.Table]:
521 """Parses a chunk of records, by table.
523 Args:
524 batch: the records of the chunk, by segment level.
525 owners: the ORI of the batch header that each record follows, by
526 segment level.
527 year: the year of the master file.
529 Returns:
530 A table for each table that has records in the chunk.
532 """
533 parsed = {}
534 for level, lines in batch.items():
535 if not lines: 535 ↛ 536line 535 didn't jump to line 536 because the condition on line 535 was never true
536 continue
537 code = level.decode()
538 table = parse(lines, code, data_year = year)
539 table = table.add_column(
540 1, 'agency_ori', _constant(owners[level], len(lines)))
541 name = layouts.TABLES[code]
542 parsed[name] = table.select(schema(name).names).cast(schema(name))
543 return parsed
546def _state_codes(states: Sequence[str] | str | None) -> set[bytes]:
547 """Returns the numeric state codes of `states`, as bytes.
549 Args:
550 states: postal abbreviations or numeric state codes, or `None`.
552 Raises:
553 ValueError: if a state is not known.
555 Returns:
556 The numeric codes (such as b"07"), or an empty set for every state.
558 """
559 numbers = {v: k for k, v in codes.STATE_ABBREVIATIONS.items()}
560 numbers.update({'NB': '26', 'GM': '55'})
561 found = set()
562 for state in utilities.to_names(states):
563 text = str(state).strip().upper()
564 if text.isdigit(): 564 ↛ 565line 564 didn't jump to line 565 because the condition on line 564 was never true
565 found.add(f'{int(text):02d}'.encode())
566 elif text in numbers:
567 found.add(numbers[text].encode())
568 else:
569 message = f'{state!r} is not a state abbreviation or code'
570 raise ValueError(message)
571 return found
574def _wanted(names: Sequence[str] | str | None) -> list[str]:
575 """Returns the names of the tables to make.
577 Args:
578 names: names of tables, or `None` for all of them.
580 Raises:
581 KeyError: if a name is not a table.
583 Returns:
584 The names.
586 """
587 wanted = utilities.to_names(names) or list(layouts.SEGMENTS)
588 unknown = [n for n in wanted if n not in layouts.SEGMENTS]
589 if unknown:
590 message = f'{unknown} are not tables: use {sorted(layouts.SEGMENTS)}'
591 raise KeyError(message)
592 return list(dict.fromkeys(wanted))
595""" Converters of the text of a field to its value """
598def _to_bool(cell: bytes, field: layouts.Field) -> bool | None: # noqa: ARG001
599 """Returns the value of a "bool" field.
601 Args:
602 cell: the characters of the field.
603 field: the field, which is not used.
605 Returns:
606 `True` for "Y", `False` for "N", and `None` otherwise.
608 """
609 text = cell.strip()
610 return True if text == b'Y' else False if text == b'N' else None
613def _to_date(
614 cell: bytes,
615 field: layouts.Field) -> dt.date | None: # noqa: ARG001
616 """Returns the value of a "date" field.
618 Args:
619 cell: the characters of the field.
620 field: the field, which is not used.
622 Returns:
623 The date written as YYYYMMDD, or `None` if it is not one.
625 """
626 text = cell.strip()
627 if len(text) != 8 or not text.isdigit():
628 return None
629 try:
630 return dt.date(int(text[:4]), int(text[4:6]), int(text[6:]))
631 except ValueError:
632 return None
635def _to_int(cell: bytes, field: layouts.Field) -> int | None: # noqa: ARG001
636 """Returns the value of an "int" field.
638 Args:
639 cell: the characters of the field.
640 field: the field, which is not used.
642 Returns:
643 The whole number, or `None` if it is not one.
645 """
646 text = cell.strip()
647 return int(text) if text.isdigit() else None
650def _to_quantity(
651 cell: bytes,
652 field: layouts.Field) -> float | None: # noqa: ARG001
653 """Returns the value of a "quantity" field.
655 Args:
656 cell: the characters of the field.
657 field: the field, which is not used.
659 Returns:
660 The amount written as 9 digits and 3 digits of thousandths, or
661 `None` if it is blank or not a number.
663 """
664 whole, fraction = cell[:9].strip(), cell[9:12].strip()
665 if not whole and not fraction:
666 return None
667 if not (whole or b'0').isdigit() or not (fraction or b'0').isdigit(): 667 ↛ 668line 667 didn't jump to line 668 because the condition on line 667 was never true
668 return None
669 return int(whole or b'0') + int(fraction.ljust(3, b'0') or b'0') / 1000
672def _to_text(cell: bytes, field: layouts.Field) -> str | None: # noqa: ARG001
673 """Returns the value of a "text" or "code" field.
675 Args:
676 cell: the characters of the field.
677 field: the field, which is not used.
679 Returns:
680 The text without blanks around it, or `None` if it is blank.
682 """
683 text = cell.decode('latin-1').strip()
684 return text or None
687def _to_yes(cell: bytes, field: layouts.Field) -> bool:
688 """Returns the value of a "yes" field.
690 Args:
691 cell: the characters of the field.
692 field: the field, whose `true` is the character that means true.
694 Returns:
695 Whether the field is the character that means true.
697 """
698 return cell.strip() == field.true.encode()
701_CONVERTERS: dict[str, Callable[[bytes, layouts.Field], Any]] = {
702 'bool': _to_bool,
703 'code': _to_text,
704 'date': _to_date,
705 'int': _to_int,
706 'quantity': _to_quantity,
707 'text': _to_text,
708 'yes': _to_yes}