Source code for sunpy.util.util

This module provides general utility functions.
import os
import hashlib
import inspect
import textwrap
from io import BytesIO
from base64 import b64encode
from shutil import get_terminal_size
from itertools import chain, count
from collections import UserList
from import Iterator

__all__ = ['unique', 'replacement_filename', 'expand_list',
           'expand_list_generator', 'dict_keys_same', 'hash_file', 'get_width',
           'get_keywords', 'get_set_methods', 'fix_duplicate_notes']

[docs] def unique(itr, key=None): """ Return only unique elements of a sequence. Parameters ---------- itr : `iterable` Any iterable sequence. key : `function`, optional A function to apply to each element in the iterable. Defaults to `None`. Returns ------- `set`: A `set` of each unique element. """ items = set() if key is None: for elem in itr: if elem not in items: yield elem items.add(elem) else: for elem in itr: x = key(elem) if x not in items: yield elem items.add(x)
[docs] def replacement_filename(path: str): """ Return a replacement path if input path is currently in use. Enumerates until an unused filename is found, e.g., "foo.fits" becomes "foo.0.fits", if that is used, "foo.1.fits" and so on. Parameters ---------- path : `str` A string path. Returns ------- `str`: A string path. """ if not os.path.exists(path): return path else: dir_, filename = os.path.split(path) base, ext = os.path.splitext(filename) for c in count(): name = base + '.' + str(c) + ext newpath = os.path.join(dir_, name) if not os.path.exists(newpath): return newpath
[docs] def expand_list(inp): """ Expand a list of lists or tuples. Parameters ---------- inp : `list`, `tuple`, `collections.UserList` The iterable to expand. Returns ------- `list` A flat list consisting of the entries of the input. References ---------- * """ return list(expand_list_generator(inp))
[docs] def expand_list_generator(inp): for item in inp: # parfive.Results are UserList if isinstance(item, (list | tuple | UserList | Iterator)): yield from expand_list_generator(item) else: yield item
def partial_key_match(key, dictionary): """ Return value/values from a dictionary based on a partial key. Each element of the partial key is matched against the keys of the dictionary and if a partial match is found the value of the key is returned. Even a partial match works here i.e even if the key matches partially a value is returned. Parameters ---------- key : `tuple` A tuple containing the partial key. dictionary : `dict` The target dictionary from which we want to retrieve the value based on the partial key. Yields ------ value The value of the matched key. References ---------- * Examples -------- >>> d = {('abc','def','ghi') : 1, ('abc', 'def', 'xyz') : 2, ('pqr', 'lmn', 'tuv') : 3} >>> list(partial_key_match(('abc', 'def', None), d)) [1, 2] """ for k, v in dictionary.items(): if all(k1 == k2 or k2 is None for k1, k2 in zip(k, key)): yield v
[docs] def dict_keys_same(list_of_dicts): """ Makes sure that a list of dictionaries all have the same keys. If a key is missing, it will be added but with a value of None. Parameters ---------- list_of_dicts : `list` of `dict` A list containing each dictionary to parse. Returns ------ `list` The list with each dict updated. References ---------- * Examples -------- >>> l = [{'x': 42}, {'x': 23, 'y': 5}] >>> dict_keys_same(l) [{'x': 42, 'y': None}, {'x': 23, 'y': 5}] """ all_keys = set(chain.from_iterable(list_of_dicts)) for d in list_of_dicts: d.update({key: None for key in all_keys if key not in d}) return list_of_dicts
[docs] def hash_file(path): """ Returns the SHA-256 hash of a file. Parameters ---------- path : `str` The path of the file to be hashed. Returns ------- `str` SHA-256 hash of the file. References ---------- * """ BUF_SIZE = 65536 sha256 = hashlib.sha256() with open(path, 'rb') as f: while True: data = if not data: break sha256.update(data) return sha256.hexdigest()
[docs] def get_width(): """ Gets the width of the current terminal. Accounts for if the 'COLUMNS' environmental variable is set. Returns ------- `int` Width of the terminal you are in. Works for IPython notebooks and normal terminals. """ width = os.environ.get("COLUMNS", None) if width: width = int(width) else: width, _ = get_terminal_size() return width
[docs] def get_keywords(func): """ Returns a set of keyword names from ``func``'s signature. Recursive if ``func`` is a list of functions and methods. Parameters ---------- func : function or method or `list` Function or method (or list of those) to extract a set of accepted keyword arguments for. Returns ------- keywords : `set` A set of accepted keyword arguments. """ if isinstance(func, list): keywords = set() for f in func: keywords.update(get_keywords(f)) return keywords sig = inspect.signature(func) keywords = { for param in sig.parameters.values() if param.default is not inspect.Parameter.empty} return keywords
[docs] def get_set_methods(obj): """ Returns a set of keyword names that can be handled by an object's ``set_...`` methods. Parameters ---------- obj : `object` Matplotlib object such as `~matplotlib.image.AxesImage` to extract handled keyword arguments for. Returns ------- keywords : `set` A set of accepted keyword arguments. Notes ----- See :meth:`matplotlib.artist.Artist.update` for an example of Matplotlib relying on this capability. """ return { m[4:] for m in dir(obj) if m.startswith('set_') and callable(getattr(obj, m, None)) }
def _figure_to_base64(fig): # Converts a matplotlib Figure to a base64 UTF-8 string buf = BytesIO() fig.savefig(buf, format='png', facecolor='none') # works better than transparent=True return b64encode(buf.getvalue()).decode('utf-8')
[docs] def fix_duplicate_notes(subclass_doc, cls_doc): """ Returns a new documentation string such that there are notes section duplication in in `` subclasses. Parameters ---------- subclass_doc : str The documentation that needs to be appended. cls_doc The original class's documentation. Returns ------- str Updated documentation that contains no note section duplication. """ existing_notes_pos = cls_doc.find('Notes\n -----') subclass_notes_pos = subclass_doc.find('Notes\n-----') subclass_notes_data = textwrap.indent(subclass_doc[subclass_notes_pos + len('Notes\n-----'):].strip(), " ") references_pattern = "References\n ----------" examples_pattern = "Examples\n -------" start_index = cls_doc.find(references_pattern if references_pattern in cls_doc else examples_pattern) if start_index!=-1: next_pattern_pos = min(pos for pos in [cls_doc.find(references_pattern, start_index), cls_doc.find(examples_pattern, start_index)] if pos != -1) other_patterns = cls_doc[:next_pattern_pos] if existing_notes_pos!=-1: cls_doc = other_patterns + subclass_notes_data.lstrip() + '\n\n ' + cls_doc[next_pattern_pos:] else: cls_doc = other_patterns + 'Notes\n -----\n' + subclass_notes_data + '\n\n ' + cls_doc[next_pattern_pos:] elif existing_notes_pos != -1: cls_doc +="\n"+subclass_notes_data else: cls_doc += textwrap.indent(subclass_doc, " ") return cls_doc