-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathutils.py
More file actions
92 lines (78 loc) · 4.35 KB
/
Copy pathutils.py
File metadata and controls
92 lines (78 loc) · 4.35 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
# -*- coding: UTF-8 -*-
import numpy as np
from math import log2
from typing import Optional
__all__ = ["ensure_str", "human_readable_size", "ngrams_counts", "ngrams_distribution", "shannon_entropy"]
shannon_entropy = lambda b: -sum([p*log2(p) for p in [float(ctr)/len(b) for ctr in [b.count(c) for c in set(b)]]]) or 0.
def ensure_str(s: str | bytes, encoding: str = "utf-8", errors: str = "strict") -> str:
""" Ensure that an input string is decoded. """
if isinstance(s, bytes):
try:
return s.decode(encoding, errors)
except:
return s.decode("latin-1")
elif not isinstance(s, (str, bytes)):
raise TypeError("not expecting type '%s'" % type(s))
return s
def human_readable_size(size: int, precision: int = 0) -> str:
""" Display bytes' size in a human-readable format given a precision. """
i, units = 0, ["B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"]
while size >= 1024 and i < len(units)-1:
i += 1
size /= 1024.0
return "%.*f%s" % (precision, size, units[i])
def ngrams_counts(byte_obj: bytes | object, n: int = 1, step: int = 1) -> list[tuple[bytes, int]]:
""" Output a sorted list of tuples (n-gram, counts) for an input byte sequence or byte object.
If the input is a byte object, the result is cached.
:param byte_obj: byte sequence ('bytes') or byte object with "bytes" and "size" attributes (i.e. pathlib2.Path)
:param n: n determining the size of n-grams, defaults to 1
:param step: step for sliding the n-grams
"""
if n not in (1, 2, 3):
raise ValueError("n must be 1, 2, or 3")
if step <= 0:
raise ValueError("step must be positive")
try:
return byte_obj._ngram_counts_cache[n]
except (AttributeError, KeyError):
pass
if isinstance(byte_obj, bytes) or hasattr(byte_obj, "bytes"):
a = np.frombuffer(data := byte_obj if isinstance(byte_obj, bytes) else byte_obj.bytes, dtype=np.uint8)
l = a.size
if l < n:
return []
if n == 1:
counts = {b.to_bytes(1, "big"): int(c) for b, c in \
enumerate(np.bincount(np.frombuffer(data, dtype=np.uint8)))}
else:
end = (m := (l - n) // step + 1) * step
grams = np.stack((a[0:end:step], a[1:1+end:step]), axis=1) if n == 2 else \
np.stack((a[0:end:step], a[1:1+end:step], a[2:2+end:step]), axis=1)
counts = {bytes(row): int(c) for row, c in zip(*np.unique(grams, axis=0, return_counts=True))}
counts = sorted(counts.items(), key=lambda p: p[1], reverse=True)
if isinstance(byte_obj, bytes):
return counts
elif hasattr(byte_obj, "bytes"):
if not hasattr(byte_obj, "_ngram_counts_cache"):
byte_obj._ngram_counts_cache = {}
if n not in byte_obj._ngram_counts_cache.keys():
byte_obj._ngram_counts_cache[n] = counts
return byte_obj._ngram_counts_cache[n]
raise TypeError(f"Bad input type ; should be a byte sequence or object (got '{type(byte_obj)}')")
def ngrams_distribution(byte_obj: bytes | object, n: int = 1, step: int = 1, n_most_common: Optional[int] = None,
n_exclude_top: int = 0, exclude: Optional[list] = None) -> list[tuple[bytes, int]]:
""" Compute the n-grams distribution of an input byte sequence or byte object given exclusions.
:param byte_obj: byte sequence ('bytes') or byte object with "bytes" and "size" attributes (i.e. pathlib2.Path)
:param n: n determining the size of n-grams, defaults to 1
:param step: step for sliding the n-grams
:param n_most_common: number of n-grams to be kept in the result, keep all by default
:param n_exclude_top: number of n-grams to be excluded from the top of the histogram, no exclusion by default
:param exclude: list of specific n-grams to be excluded, no exclusion by default
:return: list of n_most_common (n-gram, count) pairs
"""
if len(c := ngrams_counts(byte_obj, n, step)) == 0:
return []
r = c[:len(c) if n_most_common is None else n_most_common + n_exclude_top + len(exclude or [])]
if exclude is not None:
r = [(ngram, count) for ngram, count in r if ngram not in exclude]
return r[n_exclude_top:n_exclude_top+(n_most_common or len(c))]