prosodic.utils

  1from typing import Any, Callable, Dict, List, Optional, Tuple, Union
  2from .imports import *
  3
  4class SimpleCache:
  5    """A simple file-based caching system.
  6
  7    This class provides a dictionary-like interface for caching objects to disk.
  8    It uses a two-level directory structure to organize cached files.
  9
 10    Attributes:
 11        root_dir (str): The root directory for storing cached files.
 12    """
 13
 14    def __init__(self, root_dir: str = PATH_HOME_DATA_CACHE) -> None:
 15        """Initialize the SimpleCache.
 16
 17        Args:
 18            root_dir: The root directory for storing cached files.
 19        """
 20        self.root_dir = root_dir
 21        os.makedirs(root_dir, exist_ok=True)
 22
 23    def _get_file_path(self, key: str) -> str:
 24        """Get the file path for a given key.
 25
 26        Args:
 27            key: The cache key.
 28
 29        Returns:
 30            The file path for the given key.
 31        """
 32        # Use the first 2 characters for the first level directory
 33        # and the next 2 characters for the second level directory
 34        dir1 = key[:2]
 35        dir2 = key[2:4]
 36        file_name = key[4:]
 37        
 38        dir_path = os.path.join(self.root_dir, dir1, dir2)
 39        os.makedirs(dir_path, exist_ok=True)
 40        
 41        return os.path.join(dir_path, file_name)
 42
 43    def __setitem__(self, key: str, value: Any) -> None:
 44        """Set an item in the cache.
 45
 46        Args:
 47            key: The cache key.
 48            value: The value to cache.
 49        """
 50        file_path = self._get_file_path(key)
 51        with open(file_path, 'wb') as f:
 52            f.write(encode_cache(value))
 53
 54    def __getitem__(self, key: str) -> Any:
 55        """Get an item from the cache.
 56
 57        Args:
 58            key: The cache key.
 59
 60        Returns:
 61            The cached value.
 62
 63        Raises:
 64            KeyError: If the key is not found in the cache.
 65        """
 66        file_path = self._get_file_path(key)
 67        if not os.path.exists(file_path):
 68            raise KeyError(key)
 69        with open(file_path, 'rb') as f:
 70            return decode_cache(f.read())
 71
 72    def __contains__(self, key: str) -> bool:
 73        """Check if a key exists in the cache.
 74
 75        Args:
 76            key: The cache key.
 77
 78        Returns:
 79            True if the key exists, False otherwise.
 80        """
 81        return os.path.exists(self._get_file_path(key))
 82
 83    def get(self, key: str, default: Any = None) -> Any:
 84        """Get an item from the cache with a default value.
 85
 86        Args:
 87            key: The cache key.
 88            default: The default value to return if the key is not found.
 89
 90        Returns:
 91            The cached value or the default value.
 92        """
 93        try:
 94            return self[key]
 95        except KeyError:
 96            return default
 97
 98def retry_on_io_error(max_attempts: int = 3, delay: float = 0.1) -> Callable:
 99    """Decorator to retry a function on IOError.
100
101    Args:
102        max_attempts: Maximum number of retry attempts.
103        delay: Delay between retry attempts in seconds.
104
105    Returns:
106        A decorator function.
107    """
108    def decorator(func):
109        @wraps(func)
110        def wrapper(*args, **kwargs):
111            for attempt in range(max_attempts):
112                try:
113                    return func(*args, **kwargs)
114                except IOError as e:
115                    if attempt < max_attempts - 1:
116                        time.sleep(delay)
117                    else:
118                        raise
119        return wrapper
120    return decorator
121
122
123def group_ents(l: List[Any], feat: str) -> List[List[Any]]:
124    """Group entities based on a common feature.
125
126    Args:
127        l: List of entities to group.
128        feat: The feature to group by.
129
130    Returns:
131        A list of grouped entities.
132    """
133    val = None
134    ol = []
135    lx = []
136    for x in l:
137        valx = getattr(x, feat)
138        if valx is not val and lx:
139            ol.append(lx)
140            lx = []
141        lx.append(x)
142        val = valx
143    if lx:
144        ol.append(lx)
145    return ol
146
147
148def groupby(df: pd.DataFrame, groupby: Union[str, List[str]]) -> pd.core.groupby.DataFrameGroupBy:
149    """Group a DataFrame by specified columns.
150
151    Args:
152        df: The DataFrame to group.
153        groupby: Column name(s) to group by.
154
155    Returns:
156        A grouped DataFrame.
157
158    Raises:
159        Exception: If no valid grouping columns are found.
160    """
161    allcols = set(df.index.names) | {df.index.name} | set(df.columns)
162    if type(groupby) == str:
163        groupby = [groupby]
164    gby = [g for g in groupby if g in allcols]
165    if not gby:
166        raise Exception("No group after filter")
167    return df.groupby(gby)
168
169
170def get_txt(txt: Optional[str], fn: Optional[str]) -> str:
171    """Get text content from a string, file, or URL.
172
173    Args:
174        txt: Text content or None.
175        fn: Filename or URL or None.
176
177    Returns:
178        The text content.
179    """
180    if txt:
181        if txt.startswith("http") or os.path.exists(txt):
182            return get_txt(None, txt)
183
184        return txt
185
186    if fn:
187        if fn.startswith("http"):
188            response = requests.get(fn)
189            return response.text.strip()
190
191        if os.path.exists(fn):
192            with open(fn, encoding='utf-8') as f:
193                return f.read()
194
195    return ""
196
197
198def clean_text(txt: str) -> str:
199    """Clean and normalize text.
200
201    Args:
202        txt: The input text.
203
204    Returns:
205        Cleaned and normalized text.
206    """
207    txt = txt.replace("\r\n", "\n").replace("\r", "\n")
208    txt = ftfy.fix_text(txt)
209    return txt
210
211
212def get_attr_str(attrs: Dict[str, Any], sep: str = ", ", bad_keys: Optional[List[str]] = None) -> str:
213    """Generate a string representation of attributes.
214
215    Args:
216        attrs: Dictionary of attributes.
217        sep: Separator between attribute strings.
218        bad_keys: List of keys to exclude.
219
220    Returns:
221        A string representation of the attributes.
222    """
223    strs = [
224        f"{k}={repr(v)}"
225        for k, v in attrs.items()
226        if v is not None and (not bad_keys or not k in set(bad_keys))
227    ]
228    attrstr = sep.join(strs)
229    return attrstr
230
231
232def safesum(l: List[Union[int, float]]) -> Union[int, float]:
233    """Safely sum a list of numbers, ignoring non-numeric values.
234
235    Args:
236        l: List of numbers to sum.
237
238    Returns:
239        The sum of the numeric values in the list.
240    """
241    l = [x for x in l if type(x) in {int, float, np.float64, np.float32}]
242    return sum(l)
243
244
245def setindex(df: pd.DataFrame, cols: List[str] = []) -> pd.DataFrame:
246    """Set the index of a DataFrame to specified columns.
247
248    Args:
249        df: The input DataFrame.
250        cols: List of column names to set as index.
251
252    Returns:
253        The DataFrame with the new index set.
254    """
255    if not cols:
256        return df
257    cols = [c for c in cols if c in set(df.columns)]
258    return df.set_index(cols) if cols else df
259
260
261
262def get_stress(ipa: str) -> str:
263    """Get the stress level from an IPA string.
264
265    Args:
266        ipa: The IPA string.
267
268    Returns:
269        The stress level ('S', 'P', or 'U').
270    """
271    if not ipa:
272        return ""
273    if ipa[0] == "`":
274        return "S"
275    if ipa[0] == "'":
276        return "P"
277    return "U"
278
279
280def get_initial_whitespace(xstr: str) -> str:
281    """Get the initial whitespace from a string.
282
283    Args:
284        xstr: The input string.
285
286    Returns:
287        The initial whitespace.
288    """
289    o = []
290    for i, x in enumerate(xstr):
291        if x == x.strip():
292            break
293        o.append(x)
294    return "".join(o)
295
296
297def unique(l: List[Any]) -> List[Any]:
298    """Get unique elements from a list while preserving order.
299
300    Args:
301        l: The input list.
302
303    Returns:
304        A list of unique elements.
305    """
306    from ordered_set import OrderedSet
307
308    return list(OrderedSet(l))
309
310
311def hashstr(*inputs: Any, length: int = HASHSTR_LEN) -> str:
312    """Generate a hash string from inputs.
313
314    Args:
315        *inputs: Input values to hash.
316        length: Length of the output hash string.
317
318    Returns:
319        A hash string.
320    """
321    import hashlib
322
323    input_string = str(inputs)
324    sha256_hash = hashlib.sha256(str(input_string).encode()).hexdigest()
325    return sha256_hash[:length]
326
327
328def read_json(fn: str) -> Dict[str, Any]:
329    """Read a JSON file.
330
331    Args:
332        fn: The filename.
333
334    Returns:
335        The parsed JSON data as a dictionary.
336    """
337    if not os.path.exists(fn):
338        return {}
339    with open(fn, encoding='utf-8') as f:
340        return orjson.loads(f.read())
341
342
343def from_json(json_d: Union[str, Dict[str, Any]], **kwargs: Any) -> Any:
344    """Create an object from JSON data.
345
346    Args:
347        json_d: JSON data or filename.
348        **kwargs: Additional keyword arguments.
349
350    Returns:
351        The created object.
352
353    Raises:
354        Exception: If the JSON data doesn't contain a '_class' key.
355    """
356    from .imports import GLOBALS
357
358    if type(json_d) == str:
359        json_d = read_json(json_d)
360    if not "_class" in json_d:
361        pprint(json_d)
362        raise Exception
363    classname = json_d["_class"]
364    classx = GLOBALS[classname]
365    return classx.from_json(json_d, **kwargs)
366
367
368def load(fn: str, **kwargs: Any) -> Any:
369    """Load an object from a JSON file.
370
371    Args:
372        fn: The filename.
373        **kwargs: Additional keyword arguments.
374
375    Returns:
376        The loaded object.
377    """
378    return from_json(fn, **kwargs)
379
380
381def to_json(obj: Any, fn: Optional[str] = None) -> Optional[Dict[str, Any]]:
382    """Convert an object to JSON and optionally save to a file.
383
384    Args:
385        obj: The object to convert.
386        fn: The filename to save to (optional).
387
388    Returns:
389        The JSON data if fn is None, otherwise None.
390    """
391    if hasattr(obj, "to_json"):
392        data = obj.to_json()
393    else:
394        data = obj
395
396    if not fn:
397        return data
398    else:
399        fdir = os.path.dirname(fn)
400        if fdir:
401            os.makedirs(fdir, exist_ok=True)
402        with open(fn, "wb") as of:
403            of.write(
404                orjson.dumps(
405                    data, option=orjson.OPT_INDENT_2 | orjson.OPT_SERIALIZE_NUMPY
406                )
407            )
408
409
410def ensure_dir(fn: str) -> None:
411    """Ensure that the directory for a file exists.
412
413    Args:
414        fn: The filename.
415    """
416    dirname = os.path.dirname(fn)
417    if dirname:
418        os.makedirs(dirname, exist_ok=True)
419
420
421
422def encode_cache(x: Any) -> bytes:
423    """Encode an object for caching.
424
425    Args:
426        x: The object to encode.
427
428    Returns:
429        The encoded object as bytes.
430    """
431    return b64encode(
432        zlib.compress(
433            orjson.dumps(
434                x,
435                option=orjson.OPT_SERIALIZE_NUMPY,
436            )
437        )
438    )
439
440
441def decode_cache(x: bytes) -> Any:
442    """Decode a cached object.
443
444    Args:
445        x: The encoded object.
446
447    Returns:
448        The decoded object.
449    """
450    return orjson.loads(
451        zlib.decompress(
452            b64decode(
453                x,
454            ),
455        ),
456    )
457
458
459def to_html(html: Union[str, Any], as_str: bool = False, **kwargs: Any) -> Union[str, Any]:
460    """Convert an object to HTML.
461
462    Args:
463        html: The object to convert.
464        as_str: Whether to return as a string.
465        **kwargs: Additional keyword arguments.
466
467    Returns:
468        The HTML representation of the object.
469    """
470    if type(html) is not str:
471        if hasattr(html, "to_html"):
472            return html.to_html(as_str=as_str, **kwargs)
473        logger.error(f"what type of data is this? {html}")
474        return
475
476    if as_str:
477        return html
478
479    try:
480        from IPython.display import HTML, Markdown, display
481
482        return HTML(html)
483    except ModuleNotFoundError:
484        return html
485
486
487def enable_caching() -> None:
488    """Enable caching."""
489    global USE_CACHE
490    USE_CACHE = True
491
492
493def caching_is_enabled() -> bool:
494    """Check if caching is enabled.
495
496    Returns:
497        True if caching is enabled, False otherwise.
498    """
499    return USE_CACHE
500
501
502def disable_caching() -> None:
503    """Disable caching."""
504    global USE_CACHE
505    USE_CACHE = False
506
507
508@contextmanager
509def caching_enabled() -> Iterator[None]:
510    """Context manager for temporarily enabling caching."""
511    was_loud = caching_is_enabled()
512    enable_caching()
513    yield
514    if not was_loud:
515        disable_caching()
516
517
518@contextmanager
519def caching_disabled() -> Iterator[None]:
520    """Context manager for temporarily disabling caching."""
521    was_loud = caching_is_enabled()
522    disable_caching()
523    yield
524    if was_loud:
525        enable_caching()
526
527
528@contextmanager
529def logging_disabled() -> Iterator[None]:
530    """Context manager for temporarily disabling logging."""
531    was_quiet = logmap.is_quiet
532    logmap.is_quiet = True
533    yield
534    logmap.is_quiet = was_quiet
535
536
537@contextmanager
538def logging_enabled() -> Iterator[None]:
539    """Context manager for temporarily enabling logging."""
540    was_quiet = logmap.is_quiet
541    logmap.is_quiet = False
542    yield
543    logmap.is_quiet = was_quiet
544
545
546def force_int(x: Any, errors: int = 0) -> int:
547    """Convert the input to an integer.
548
549    Args:
550        x: The input value to be converted to an integer.
551        errors: The value to be returned in case of an error. Defaults to 0.
552
553    Returns:
554        The input value converted to an integer if successful, otherwise the specified error value.
555    """
556    try:
557        return int(x)
558    except (ValueError, TypeError):
559        return errors
560
561def tokenize_agnostic(txt: str) -> List[str]:
562    """Tokenize text in a language-agnostic way.
563
564    Args:
565        txt: The input text.
566
567    Returns:
568        A list of tokens.
569    """
570    return re.findall(r"[\w']+|[.,!?; -—–'\n]", txt)