src/pyutils/files/file_utils.py

   1 #!/usr/bin/env python3
   2
   3 # © Copyright 2021-2022, Scott Gasch
   4
   5 """Utilities for working with files."""
   6
   7 import contextlib
   8 import datetime
   9 import errno
  10 import fnmatch
  11 import glob
  12 import hashlib
  13 import logging
  14 import os
  15 import pathlib
  16 import re
  17 import time
  18 from os.path import exists, isfile, join
  19 from typing import Callable, List, Literal, Optional, TextIO
  20 from uuid import uuid4
  21
  22 logger = logging.getLogger(__name__)
  23
  24
  25 def remove_newlines(x: str) -> str:
  26     """Trivial function to be used as a line_transformer in
  27     :meth:`slurp_file` for no newlines in file contents"""
  28     return x.replace('\n', '')
  29
  30
  31 def strip_whitespace(x: str) -> str:
  32     """Trivial function to be used as a line_transformer in
  33     :meth:`slurp_file` for no leading / trailing whitespace in
  34     file contents"""
  35     return x.strip()
  36
  37
  38 def remove_hash_comments(x: str) -> str:
  39     """Trivial function to be used as a line_transformer in
  40     :meth:`slurp_file` for no # comments in file contents"""
  41     return re.sub(r'#.*$', '', x)
  42
  43
  44 def slurp_file(
  45     filename: str,
  46     *,
  47     skip_blank_lines=False,
  48     line_transformers: Optional[List[Callable[[str], str]]] = None,
  49 ):
  50     """Reads in a file's contents line-by-line to a memory buffer applying
  51     each line transformation in turn.
  52
  53     Args:
  54         filename: file to be read
  55         skip_blank_lines: should reading skip blank lines?
  56         line_transformers: little string->string transformations
  57     """
  58
  59     ret = []
  60     xforms = []
  61     if line_transformers is not None:
  62         for x in line_transformers:
  63             xforms.append(x)
  64     if not file_is_readable(filename):
  65         raise Exception(f'{filename} can\'t be read.')
  66     with open(filename) as rf:
  67         for line in rf:
  68             for transformation in xforms:
  69                 line = transformation(line)
  70             if skip_blank_lines and line == '':
  71                 continue
  72             ret.append(line)
  73     return ret
  74
  75
  76 def remove(path: str) -> None:
  77     """Deletes a file.  Raises if path refers to a directory or a file
  78     that doesn't exist.
  79
  80     Args:
  81         path: the path of the file to delete
  82
  83     >>> import os
  84     >>> filename = '/tmp/file_utils_test_file'
  85     >>> os.system(f'touch {filename}')
  86     0
  87     >>> does_file_exist(filename)
  88     True
  89     >>> remove(filename)
  90     >>> does_file_exist(filename)
  91     False
  92     """
  93     os.remove(path)
  94
  95
  96 def fix_multiple_slashes(path: str) -> str:
  97     """Fixes multi-slashes in paths or path-like strings
  98
  99     Args:
 100         path: the path in which to remove multiple slashes
 101
 102     >>> p = '/usr/local//etc/rc.d///file.txt'
 103     >>> fix_multiple_slashes(p)
 104     '/usr/local/etc/rc.d/file.txt'
 105
 106     >>> p = 'this is a test'
 107     >>> fix_multiple_slashes(p) == p
 108     True
 109     """
 110     return re.sub(r'/+', '/', path)
 111
 112
 113 def delete(path: str) -> None:
 114     """This is a convenience for my dumb ass who can't remember os.remove
 115     sometimes.
 116     """
 117     os.remove(path)
 118
 119
 120 def without_extension(path: str) -> str:
 121     """Remove one (the last) extension from a file or path.
 122
 123     Args:
 124         path: the path from which to remove an extension
 125
 126     Returns:
 127         the path with one extension removed.
 128
 129     >>> without_extension('foobar.txt')
 130     'foobar'
 131
 132     >>> without_extension('/home/scott/frapp.py')
 133     '/home/scott/frapp'
 134
 135     >>> f = 'a.b.c.tar.gz'
 136     >>> while('.' in f):
 137     ...     f = without_extension(f)
 138     ...     print(f)
 139     a.b.c.tar
 140     a.b.c
 141     a.b
 142     a
 143
 144     >>> without_extension('foobar')
 145     'foobar'
 146
 147     """
 148     return os.path.splitext(path)[0]
 149
 150
 151 def without_all_extensions(path: str) -> str:
 152     """Removes all extensions from a path; handles multiple extensions
 153     like foobar.tar.gz -> foobar.
 154
 155     Args:
 156         path: the path from which to remove all extensions
 157
 158     Returns:
 159         the path with all extensions removed.
 160
 161     >>> without_all_extensions('/home/scott/foobar.1.tar.gz')
 162     '/home/scott/foobar'
 163
 164     """
 165     while '.' in path:
 166         path = without_extension(path)
 167     return path
 168
 169
 170 def get_extension(path: str) -> str:
 171     """Extract and return one (the last) extension from a file or path.
 172
 173     Args:
 174         path: the path from which to extract an extension
 175
 176     Returns:
 177         The last extension from the file path.
 178
 179     >>> get_extension('this_is_a_test.txt')
 180     '.txt'
 181
 182     >>> get_extension('/home/scott/test.py')
 183     '.py'
 184
 185     >>> get_extension('foobar')
 186     ''
 187
 188     """
 189     return os.path.splitext(path)[1]
 190
 191
 192 def get_all_extensions(path: str) -> List[str]:
 193     """Return the extensions of a file or path in order.
 194
 195     Args:
 196         path: the path from which to extract all extensions.
 197
 198     Returns:
 199         a list containing each extension which may be empty.
 200
 201     >>> get_all_extensions('/home/scott/foo.tar.gz.1')
 202     ['.tar', '.gz', '.1']
 203
 204     >>> get_all_extensions('/home/scott/foobar')
 205     []
 206
 207     """
 208     ret = []
 209     while True:
 210         ext = get_extension(path)
 211         path = without_extension(path)
 212         if ext:
 213             ret.append(ext)
 214         else:
 215             ret.reverse()
 216             return ret
 217
 218
 219 def without_path(filespec: str) -> str:
 220     """Returns the base filename without any leading path.
 221
 222     Args:
 223         filespec: path to remove leading directories from
 224
 225     Returns:
 226         filespec without leading dir components.
 227
 228     >>> without_path('/home/scott/foo.py')
 229     'foo.py'
 230
 231     >>> without_path('foo.py')
 232     'foo.py'
 233
 234     """
 235     return os.path.split(filespec)[1]
 236
 237
 238 def get_path(filespec: str) -> str:
 239     """Returns just the path of the filespec by removing the filename and
 240     extension.
 241
 242     Args:
 243         filespec: path to remove filename / extension(s) from
 244
 245     Returns:
 246         filespec with just the leading directory components and no
 247             filename or extension(s)
 248
 249     >>> get_path('/home/scott/foobar.py')
 250     '/home/scott'
 251
 252     >>> get_path('/home/scott/test.1.2.3.gz')
 253     '/home/scott'
 254
 255     >>> get_path('~scott/frapp.txt')
 256     '~scott'
 257
 258     """
 259     return os.path.split(filespec)[0]
 260
 261
 262 def get_canonical_path(filespec: str) -> str:
 263     """Returns a canonicalized absolute path.
 264
 265     Args:
 266         filespec: the path to canonicalize
 267
 268     Returns:
 269         the canonicalized path
 270
 271     >>> get_canonical_path('/home/scott/../../home/lynn/../scott/foo.txt')
 272     '/usr/home/scott/foo.txt'
 273
 274     """
 275     return os.path.realpath(filespec)
 276
 277
 278 def create_path_if_not_exist(path, on_error=None) -> None:
 279     """
 280     Attempts to create path if it does not exist already.
 281
 282     .. warning::
 283
 284         Files are created with mode 0x0777 (i.e. world read/writeable).
 285
 286     Args:
 287         path: the path to attempt to create
 288         on_error: If True, it's invoked on error conditions.  Otherwise
 289             any exceptions are raised.
 290
 291     >>> import uuid
 292     >>> import os
 293     >>> path = os.path.join("/tmp", str(uuid.uuid4()), str(uuid.uuid4()))
 294     >>> os.path.exists(path)
 295     False
 296     >>> create_path_if_not_exist(path)
 297     >>> os.path.exists(path)
 298     True
 299     """
 300     logger.debug("Creating path %s", path)
 301     previous_umask = os.umask(0)
 302     try:
 303         os.makedirs(path)
 304         os.chmod(path, 0o777)
 305     except OSError as ex:
 306         if ex.errno != errno.EEXIST and not os.path.isdir(path):
 307             if on_error is not None:
 308                 on_error(path, ex)
 309             else:
 310                 raise
 311     finally:
 312         os.umask(previous_umask)
 313
 314
 315 def does_file_exist(filename: str) -> bool:
 316     """Returns True if a file exists and is a normal file.
 317
 318     Args:
 319         filename: filename to check
 320
 321     Returns:
 322         True if filename exists and is a normal file.
 323
 324     >>> does_file_exist(__file__)
 325     True
 326     >>> does_file_exist('/tmp/2492043r9203r9230r9230r49230r42390r4230')
 327     False
 328     """
 329     return os.path.exists(filename) and os.path.isfile(filename)
 330
 331
 332 def file_is_readable(filename: str) -> bool:
 333     """True if file exists, is a normal file and is readable by the
 334     current process.  False otherwise.
 335
 336     Args:
 337         filename: the filename to check for read access
 338     """
 339     return does_file_exist(filename) and os.access(filename, os.R_OK)
 340
 341
 342 def file_is_writable(filename: str) -> bool:
 343     """True if file exists, is a normal file and is writable by the
 344     current process.  False otherwise.
 345
 346     Args:
 347         filename: the file to check for write access.
 348     """
 349     return does_file_exist(filename) and os.access(filename, os.W_OK)
 350
 351
 352 def file_is_executable(filename: str) -> bool:
 353     """True if file exists, is a normal file and is executable by the
 354     current process.  False otherwise.
 355
 356     Args:
 357         filename: the file to check for execute access.
 358     """
 359     return does_file_exist(filename) and os.access(filename, os.X_OK)
 360
 361
 362 def does_directory_exist(dirname: str) -> bool:
 363     """Returns True if a file exists and is a directory.
 364
 365     >>> does_directory_exist('/tmp')
 366     True
 367     >>> does_directory_exist('/xyzq/21341')
 368     False
 369     """
 370     return os.path.exists(dirname) and os.path.isdir(dirname)
 371
 372
 373 def does_path_exist(pathname: str) -> bool:
 374     """Just a more verbose wrapper around os.path.exists."""
 375     return os.path.exists(pathname)
 376
 377
 378 def get_file_size(filename: str) -> int:
 379     """Returns the size of a file in bytes.
 380
 381     Args:
 382         filename: the filename to size
 383
 384     Returns:
 385         size of filename in bytes
 386     """
 387     return os.path.getsize(filename)
 388
 389
 390 def is_normal_file(filename: str) -> bool:
 391     """Returns True if filename is a normal file.
 392
 393     >>> is_normal_file(__file__)
 394     True
 395     """
 396     return os.path.isfile(filename)
 397
 398
 399 def is_directory(filename: str) -> bool:
 400     """Returns True if filename is a directory.
 401
 402     >>> is_directory('/tmp')
 403     True
 404     """
 405     return os.path.isdir(filename)
 406
 407
 408 def is_symlink(filename: str) -> bool:
 409     """True if filename is a symlink, False otherwise.
 410
 411     >>> is_symlink('/tmp')
 412     False
 413
 414     >>> is_symlink('/home')
 415     True
 416
 417     """
 418     return os.path.islink(filename)
 419
 420
 421 def is_same_file(file1: str, file2: str) -> bool:
 422     """Returns True if the two files are the same inode.
 423
 424     >>> is_same_file('/tmp', '/tmp/../tmp')
 425     True
 426
 427     >>> is_same_file('/tmp', '/home')
 428     False
 429
 430     """
 431     return os.path.samefile(file1, file2)
 432
 433
 434 def get_file_raw_timestamps(filename: str) -> Optional[os.stat_result]:
 435     """Stats the file and returns an os.stat_result or None on error.
 436
 437     Args:
 438         filename: the file whose timestamps to fetch
 439
 440     Returns:
 441         the os.stat_result or None to indicate an error occurred
 442     """
 443     try:
 444         return os.stat(filename)
 445     except Exception as e:
 446         logger.exception(e)
 447         return None
 448
 449
 450 def get_file_raw_timestamp(
 451     filename: str, extractor: Callable[[os.stat_result], Optional[float]]
 452 ) -> Optional[float]:
 453     """Stat a file and, if successful, use extractor to fetch some
 454     subset of the information in the os.stat_result.  See also
 455     :meth:`get_file_raw_atime`, :meth:`get_file_raw_mtime`, and
 456     :meth:`get_file_raw_ctime` which just call this with a lambda
 457     extractor.
 458
 459     Args:
 460         filename: the filename to stat
 461         extractor: Callable that takes a os.stat_result and produces
 462             something useful(?) with it.
 463
 464     Returns:
 465         whatever the extractor produced or None on error.
 466     """
 467     tss = get_file_raw_timestamps(filename)
 468     if tss is not None:
 469         return extractor(tss)
 470     return None
 471
 472
 473 def get_file_raw_atime(filename: str) -> Optional[float]:
 474     """Get a file's raw access time or None on error.
 475
 476     See also :meth:`get_file_atime_as_datetime`,
 477     :meth:`get_file_atime_timedelta`,
 478     and :meth:`get_file_atime_age_seconds`.
 479     """
 480     return get_file_raw_timestamp(filename, lambda x: x.st_atime)
 481
 482
 483 def get_file_raw_mtime(filename: str) -> Optional[float]:
 484     """Get a file's raw modification time or None on error.
 485
 486     See also :meth:`get_file_mtime_as_datetime`,
 487     :meth:`get_file_mtime_timedelta`,
 488     and :meth:`get_file_mtime_age_seconds`.
 489     """
 490     return get_file_raw_timestamp(filename, lambda x: x.st_mtime)
 491
 492
 493 def get_file_raw_ctime(filename: str) -> Optional[float]:
 494     """Get a file's raw creation time or None on error.
 495
 496     See also :meth:`get_file_ctime_as_datetime`,
 497     :meth:`get_file_ctime_timedelta`,
 498     and :meth:`get_file_ctime_age_seconds`.
 499     """
 500     return get_file_raw_timestamp(filename, lambda x: x.st_ctime)
 501
 502
 503 def get_file_md5(filename: str) -> str:
 504     """Hashes filename's disk contents and returns the MD5 digest.
 505
 506     Args:
 507         filename: the file whose contents to hash
 508
 509     Returns:
 510         the MD5 digest of the file's contents.  Raises on errors.
 511     """
 512     file_hash = hashlib.md5()
 513     with open(filename, "rb") as f:
 514         chunk = f.read(8192)
 515         while chunk:
 516             file_hash.update(chunk)
 517             chunk = f.read(8192)
 518     return file_hash.hexdigest()
 519
 520
 521 def set_file_raw_atime(filename: str, atime: float):
 522     """Sets a file's raw access time.
 523
 524     See also :meth:`get_file_atime_as_datetime`,
 525     :meth:`get_file_atime_timedelta`,
 526     :meth:`get_file_atime_age_seconds`,
 527     and :meth:`get_file_raw_atime`.
 528     """
 529     mtime = get_file_raw_mtime(filename)
 530     assert mtime is not None
 531     os.utime(filename, (atime, mtime))
 532
 533
 534 def set_file_raw_mtime(filename: str, mtime: float):
 535     """Sets a file's raw modification time.
 536
 537     See also :meth:`get_file_mtime_as_datetime`,
 538     :meth:`get_file_mtime_timedelta`,
 539     :meth:`get_file_mtime_age_seconds`,
 540     and :meth:`get_file_raw_mtime`.
 541     """
 542     atime = get_file_raw_atime(filename)
 543     assert atime is not None
 544     os.utime(filename, (atime, mtime))
 545
 546
 547 def set_file_raw_atime_and_mtime(filename: str, ts: float = None):
 548     """Sets both a file's raw modification and access times
 549
 550     Args:
 551         filename: the file whose times to set
 552         ts: the raw time to set or None to indicate time should be
 553             set to the current time.
 554     """
 555     if ts is not None:
 556         os.utime(filename, (ts, ts))
 557     else:
 558         os.utime(filename, None)
 559
 560
 561 def convert_file_timestamp_to_datetime(
 562     filename: str, producer
 563 ) -> Optional[datetime.datetime]:
 564     """Convert a raw file timestamp into a python datetime."""
 565     ts = producer(filename)
 566     if ts is not None:
 567         return datetime.datetime.fromtimestamp(ts)
 568     return None
 569
 570
 571 def get_file_atime_as_datetime(filename: str) -> Optional[datetime.datetime]:
 572     """Fetch a file's access time as a python datetime.
 573
 574     See also :meth:`get_file_atime_as_datetime`,
 575     :meth:`get_file_atime_timedelta`,
 576     :meth:`get_file_atime_age_seconds`,
 577     :meth:`describe_file_atime`,
 578     and :meth:`get_file_raw_atime`.
 579     """
 580     return convert_file_timestamp_to_datetime(filename, get_file_raw_atime)
 581
 582
 583 def get_file_mtime_as_datetime(filename: str) -> Optional[datetime.datetime]:
 584     """Fetches a file's modification time as a python datetime.
 585
 586     See also :meth:`get_file_mtime_as_datetime`,
 587     :meth:`get_file_mtime_timedelta`,
 588     :meth:`get_file_mtime_age_seconds`,
 589     and :meth:`get_file_raw_mtime`.
 590     """
 591     return convert_file_timestamp_to_datetime(filename, get_file_raw_mtime)
 592
 593
 594 def get_file_ctime_as_datetime(filename: str) -> Optional[datetime.datetime]:
 595     """Fetches a file's creation time as a python datetime.
 596
 597     See also :meth:`get_file_ctime_as_datetime`,
 598     :meth:`get_file_ctime_timedelta`,
 599     :meth:`get_file_ctime_age_seconds`,
 600     and :meth:`get_file_raw_ctime`.
 601     """
 602     return convert_file_timestamp_to_datetime(filename, get_file_raw_ctime)
 603
 604
 605 def get_file_timestamp_age_seconds(filename: str, extractor) -> Optional[int]:
 606     """~Internal helper"""
 607     now = time.time()
 608     ts = get_file_raw_timestamps(filename)
 609     if ts is None:
 610         return None
 611     result = extractor(ts)
 612     return now - result
 613
 614
 615 def get_file_atime_age_seconds(filename: str) -> Optional[int]:
 616     """Gets a file's access time as an age in seconds (ago).
 617
 618     See also :meth:`get_file_atime_as_datetime`,
 619     :meth:`get_file_atime_timedelta`,
 620     :meth:`get_file_atime_age_seconds`,
 621     :meth:`describe_file_atime`,
 622     and :meth:`get_file_raw_atime`.
 623     """
 624     return get_file_timestamp_age_seconds(filename, lambda x: x.st_atime)
 625
 626
 627 def get_file_ctime_age_seconds(filename: str) -> Optional[int]:
 628     """Gets a file's creation time as an age in seconds (ago).
 629
 630     See also :meth:`get_file_ctime_as_datetime`,
 631     :meth:`get_file_ctime_timedelta`,
 632     :meth:`get_file_ctime_age_seconds`,
 633     and :meth:`get_file_raw_ctime`.
 634     """
 635     return get_file_timestamp_age_seconds(filename, lambda x: x.st_ctime)
 636
 637
 638 def get_file_mtime_age_seconds(filename: str) -> Optional[int]:
 639     """Gets a file's modification time as seconds (ago).
 640
 641     See also :meth:`get_file_mtime_as_datetime`,
 642     :meth:`get_file_mtime_timedelta`,
 643     :meth:`get_file_mtime_age_seconds`,
 644     and :meth:`get_file_raw_mtime`.
 645     """
 646     return get_file_timestamp_age_seconds(filename, lambda x: x.st_mtime)
 647
 648
 649 def get_file_timestamp_timedelta(
 650     filename: str, extractor
 651 ) -> Optional[datetime.timedelta]:
 652     """~Internal helper"""
 653     age = get_file_timestamp_age_seconds(filename, extractor)
 654     if age is not None:
 655         return datetime.timedelta(seconds=float(age))
 656     return None
 657
 658
 659 def get_file_atime_timedelta(filename: str) -> Optional[datetime.timedelta]:
 660     """How long ago was a file accessed as a timedelta?
 661
 662     See also :meth:`get_file_atime_as_datetime`,
 663     :meth:`get_file_atime_timedelta`,
 664     :meth:`get_file_atime_age_seconds`,
 665     :meth:`describe_file_atime`,
 666     and :meth:`get_file_raw_atime`.
 667     """
 668     return get_file_timestamp_timedelta(filename, lambda x: x.st_atime)
 669
 670
 671 def get_file_ctime_timedelta(filename: str) -> Optional[datetime.timedelta]:
 672     """How long ago was a file created as a timedelta?
 673
 674     See also :meth:`get_file_ctime_as_datetime`,
 675     :meth:`get_file_ctime_timedelta`,
 676     :meth:`get_file_ctime_age_seconds`,
 677     and :meth:`get_file_raw_ctime`.
 678     """
 679     return get_file_timestamp_timedelta(filename, lambda x: x.st_ctime)
 680
 681
 682 def get_file_mtime_timedelta(filename: str) -> Optional[datetime.timedelta]:
 683     """
 684     Gets a file's modification time as a python timedelta.
 685
 686     See also :meth:`get_file_mtime_as_datetime`,
 687     :meth:`get_file_mtime_timedelta`,
 688     :meth:`get_file_mtime_age_seconds`,
 689     and :meth:`get_file_raw_mtime`.
 690     """
 691     return get_file_timestamp_timedelta(filename, lambda x: x.st_mtime)
 692
 693
 694 def describe_file_timestamp(filename: str, extractor, *, brief=False) -> Optional[str]:
 695     """~Internal helper"""
 696     from pyutils.datetimez.datetime_utils import (
 697         describe_duration,
 698         describe_duration_briefly,
 699     )
 700
 701     age = get_file_timestamp_age_seconds(filename, extractor)
 702     if age is None:
 703         return None
 704     if brief:
 705         return describe_duration_briefly(age)
 706     else:
 707         return describe_duration(age)
 708
 709
 710 def describe_file_atime(filename: str, *, brief=False) -> Optional[str]:
 711     """
 712     Describe how long ago a file was accessed.
 713
 714     See also :meth:`get_file_atime_as_datetime`,
 715     :meth:`get_file_atime_timedelta`,
 716     :meth:`get_file_atime_age_seconds`,
 717     :meth:`describe_file_atime`,
 718     and :meth:`get_file_raw_atime`.
 719     """
 720     return describe_file_timestamp(filename, lambda x: x.st_atime, brief=brief)
 721
 722
 723 def describe_file_ctime(filename: str, *, brief=False) -> Optional[str]:
 724     """Describes a file's creation time.
 725
 726     See also :meth:`get_file_ctime_as_datetime`,
 727     :meth:`get_file_ctime_timedelta`,
 728     :meth:`get_file_ctime_age_seconds`,
 729     and :meth:`get_file_raw_ctime`.
 730     """
 731     return describe_file_timestamp(filename, lambda x: x.st_ctime, brief=brief)
 732
 733
 734 def describe_file_mtime(filename: str, *, brief=False) -> Optional[str]:
 735     """
 736     Describes how long ago a file was modified.
 737
 738     See also :meth:`get_file_mtime_as_datetime`,
 739     :meth:`get_file_mtime_timedelta`,
 740     :meth:`get_file_mtime_age_seconds`,
 741     and :meth:`get_file_raw_mtime`.
 742     """
 743     return describe_file_timestamp(filename, lambda x: x.st_mtime, brief=brief)
 744
 745
 746 def touch_file(filename: str, *, mode: Optional[int] = 0o666):
 747     """Like unix "touch" command's semantics: update the timestamp
 748     of a file to the current time if the file exists.  Create the
 749     file if it doesn't exist.
 750
 751     Args:
 752         filename: the filename
 753         mode: the mode to create the file with
 754     """
 755     pathlib.Path(filename, mode=mode).touch()
 756
 757
 758 def expand_globs(in_filename: str):
 759     """Expands shell globs (* and ? wildcards) to the matching files."""
 760     for filename in glob.glob(in_filename):
 761         yield filename
 762
 763
 764 def get_files(directory: str):
 765     """Returns the files in a directory as a generator."""
 766     for filename in os.listdir(directory):
 767         full_path = join(directory, filename)
 768         if isfile(full_path) and exists(full_path):
 769             yield full_path
 770
 771
 772 def get_matching_files(directory: str, glob: str):
 773     """Returns the subset of files whose name matches a glob."""
 774     for filename in get_files(directory):
 775         if fnmatch.fnmatch(filename, glob):
 776             yield filename
 777
 778
 779 def get_directories(directory: str):
 780     """Returns the subdirectories in a directory as a generator."""
 781     for d in os.listdir(directory):
 782         full_path = join(directory, d)
 783         if not isfile(full_path) and exists(full_path):
 784             yield full_path
 785
 786
 787 def get_files_recursive(directory: str):
 788     """Find the files and directories under a root recursively."""
 789     for filename in get_files(directory):
 790         yield filename
 791     for subdir in get_directories(directory):
 792         for file_or_directory in get_files_recursive(subdir):
 793             yield file_or_directory
 794
 795
 796 def get_matching_files_recursive(directory: str, glob: str):
 797     """Returns the subset of files whose name matches a glob under a root recursively."""
 798     for filename in get_files_recursive(directory):
 799         if fnmatch.fnmatch(filename, glob):
 800             yield filename
 801
 802
 803 class FileWriter(contextlib.AbstractContextManager):
 804     """A helper that writes a file to a temporary location and then moves
 805     it atomically to its ultimate destination on close.
 806     """
 807
 808     def __init__(self, filename: str) -> None:
 809         self.filename = filename
 810         uuid = uuid4()
 811         self.tempfile = f'{filename}-{uuid}.tmp'
 812         self.handle: Optional[TextIO] = None
 813
 814     def __enter__(self) -> TextIO:
 815         assert not does_path_exist(self.tempfile)
 816         self.handle = open(self.tempfile, mode="w")
 817         return self.handle
 818
 819     def __exit__(self, exc_type, exc_val, exc_tb) -> Literal[False]:
 820         if self.handle is not None:
 821             self.handle.close()
 822             cmd = f'/bin/mv -f {self.tempfile} {self.filename}'
 823             ret = os.system(cmd)
 824             if (ret >> 8) != 0:
 825                 raise Exception(f'{cmd} failed, exit value {ret>>8}!')
 826         return False
 827
 828
 829 if __name__ == '__main__':
 830     import doctest
 831
 832     doctest.testmod()