darwin/cogapp/utils.py

d775ee353a Oliv*0001 """ Utilities for cog.
                0002 """
                0003
                0004 import contextlib
                0005 import functools
                0006 import hashlib
                0007 import os
                0008 import sys
                0009
                0010
                0011 # Support FIPS mode where possible (Python >= 3.9). We don't use MD5 for security.
                0012 md5 = (
                0013     functools.partial(hashlib.md5, usedforsecurity=False)
                0014     if sys.version_info >= (3, 9)
                0015     else hashlib.md5
                0016 )
                0017
                0018
                0019 class Redirectable:
                0020     """ An object with its own stdout and stderr files.
                0021     """
                0022     def __init__(self):
                0023         self.stdout = sys.stdout
                0024         self.stderr = sys.stderr
                0025
                0026     def setOutput(self, stdout=None, stderr=None):
                0027         """ Assign new files for standard out and/or standard error.
                0028         """
                0029         if stdout:
                0030             self.stdout = stdout
                0031         if stderr:
                0032             self.stderr = stderr
                0033
                0034     def prout(self, s, end="\n"):
                0035         print(s, file=self.stdout, end=end)
                0036
                0037     def prerr(self, s, end="\n"):
                0038         print(s, file=self.stderr, end=end)
                0039
                0040
                0041 class NumberedFileReader:
                0042     """ A decorator for files that counts the readline()'s called.
                0043     """
                0044     def __init__(self, f):
                0045         self.f = f
                0046         self.n = 0
                0047
                0048     def readline(self):
                0049         l = self.f.readline()
                0050         if l:
                0051             self.n += 1
                0052         return l
                0053
                0054     def linenumber(self):
                0055         return self.n
                0056
                0057
                0058 @contextlib.contextmanager
                0059 def change_dir(new_dir):
                0060     """Change directory, and then change back.
                0061
                0062     Use as a context manager, it will return to the original
                0063     directory at the end of the block.
                0064
                0065     """
                0066     old_dir = os.getcwd()
                0067     os.chdir(str(new_dir))
                0068     try:
                0069         yield
                0070     finally:
                0071         os.chdir(old_dir)