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

__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 ---------- * https://stackoverflow.com/questions/2185822/expanding-elements-in-a-list/2185971#2185971 """ return [item for item in expand_list_generator(inp)]
[docs] def expand_list_generator(inp): for item in inp: if isinstance(item, (tuple, list, UserList)): 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 ---------- * https://stackoverflow.com/questions/18893624/partial-match-dictionary-keyof-tuples-in-python 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 ---------- * https://stackoverflow.com/questions/10482439/make-sure-all-dicts-in-a-list-have-the-same-keys 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 ---------- * https://stackoverflow.com/a/22058673 """ BUF_SIZE = 65536 sha256 = hashlib.sha256() with open(path, 'rb') as f: while True: data = f.read(BUF_SIZE) 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 = {param.name 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 `~sunpy.map.Map` 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