Coverage for databooks/common.py: 100%

1"""Common set of miscellaneous functions.""" 

2import json 

3import logging 

4import os 

5from itertools import chain 

6from pathlib import Path 

7from typing import List 

8 

9from rich.logging import RichHandler 

10 

11from databooks import JupyterNotebook 

12 

13 

14def get_logger(name: str) -> logging.Logger: 

15 """Get logger with rich configuration.""" 

16 level = os.getenv("LOG_LEVEL", logging.INFO) 

17 

18 logging.basicConfig( 

19 level=level, 

20 format="%(message)s", 

21 datefmt="[%X]", 

22 handlers=[RichHandler(rich_tracebacks=True)], 

23 ) 

24 return logging.getLogger(name) 

25 

26 

27def set_verbose(logger: logging.Logger) -> None: 

28 """Set logger to DEBUG level when user requests verbosity.""" 

29 verbose_level = logging.DEBUG 

30 logger.setLevel(verbose_level) 

31 logger.debug( 

32 f"Verbose mode: setting log level to {logging.getLevelName(verbose_level)}" 

33 ) 

34 

35 

36def write_notebook(nb: JupyterNotebook, path: Path) -> None: 

37 """Write notebook to a path.""" 

38 with path.open("w") as f: 

39 json.dump(nb.dict(), fp=f, indent=2) 

40 

41 

42def expand_paths(paths: List[Path], ignore: List[str]) -> List[Path]: 

43 """ 

44 Get paths of existing file from list of directory or file paths. 

45 

46 :param paths: Paths to consider (can be directories or files) 

47 :param ignore: Glob expressions of files to ignore 

48 :return: List of existing paths for notebooks 

49 """ 

50 paths = list( 

51 chain.from_iterable( 

52 list(path.rglob("*.ipynb")) if path.is_dir() else [path] for path in paths 

53 ) 

54 ) 

55 

56 return [ 

57 p 

58 for p in paths 

59 if not any(p.match(i) for i in ignore) and p.exists() and p.suffix == ".ipynb" 

60 ]