networkx.utils.decorators 源代码

import sys
from warnings import warn

from collections import defaultdict
from os.path import splitext
from contextlib import contextmanager
try:
    from pathlib import Path
except ImportError:
    # Use Path to indicate if pathlib exists (like numpy does)
    Path = None

import networkx as nx
from decorator import decorator
from networkx.utils import is_string_like, create_random_state, \
                           create_py_random_state

__all__ = [
    'not_implemented_for',
    'open_file',
    'nodes_or_number',
    'preserve_random_state',
    'random_state',
    'np_random_state',
    'py_random_state',
]


[文档]def not_implemented_for(*graph_types): """Decorator to mark algorithms as not implemented Parameters ---------- graph_types : container of strings Entries must be one of 'directed','undirected', 'multigraph', 'graph'. Returns ------- _require : function The decorated function. Raises ------ NetworkXNotImplemented If any of the packages cannot be imported Notes ----- Multiple types are joined logically with "and". For "or" use multiple @not_implemented_for() lines. Examples -------- Decorate functions like this:: @not_implemnted_for('directed') def sp_function(G): pass @not_implemnted_for('directed','multigraph') def sp_np_function(G): pass """ @decorator def _not_implemented_for(not_implement_for_func, *args, **kwargs): graph = args[0] terms = {'directed': graph.is_directed(), 'undirected': not graph.is_directed(), 'multigraph': graph.is_multigraph(), 'graph': not graph.is_multigraph()} match = True try: for t in graph_types: match = match and terms[t] except KeyError: raise KeyError('use one or more of ', 'directed, undirected, multigraph, graph') if match: msg = 'not implemented for %s type' % ' '.join(graph_types) raise nx.NetworkXNotImplemented(msg) else: return not_implement_for_func(*args, **kwargs) return _not_implemented_for
def _open_gz(path, mode): import gzip return gzip.open(path, mode=mode) def _open_bz2(path, mode): import bz2 return bz2.BZ2File(path, mode=mode) # To handle new extensions, define a function accepting a `path` and `mode`. # Then add the extension to _dispatch_dict. _dispatch_dict = defaultdict(lambda: open) _dispatch_dict['.gz'] = _open_gz _dispatch_dict['.bz2'] = _open_bz2 _dispatch_dict['.gzip'] = _open_gz
[文档]def open_file(path_arg, mode='r'): """Decorator to ensure clean opening and closing of files. Parameters ---------- path_arg : int Location of the path argument in args. Even if the argument is a named positional argument (with a default value), you must specify its index as a positional argument. mode : str String for opening mode. Returns ------- _open_file : function Function which cleanly executes the io. Examples -------- Decorate functions like this:: @open_file(0,'r') def read_function(pathname): pass @open_file(1,'w') def write_function(G,pathname): pass @open_file(1,'w') def write_function(G, pathname='graph.dot') pass @open_file('path', 'w+') def another_function(arg, **kwargs): path = kwargs['path'] pass """ # Note that this decorator solves the problem when a path argument is # specified as a string, but it does not handle the situation when the # function wants to accept a default of None (and then handle it). # Here is an example: # # @open_file('path') # def some_function(arg1, arg2, path=None): # if path is None: # fobj = tempfile.NamedTemporaryFile(delete=False) # close_fobj = True # else: # # `path` could have been a string or file object or something # # similar. In any event, the decorator has given us a file object # # and it will close it for us, if it should. # fobj = path # close_fobj = False # # try: # fobj.write('blah') # finally: # if close_fobj: # fobj.close() # # Normally, we'd want to use "with" to ensure that fobj gets closed. # However, recall that the decorator will make `path` a file object for # us, and using "with" would undesirably close that file object. Instead, # you use a try block, as shown above. When we exit the function, fobj will # be closed, if it should be, by the decorator. @decorator def _open_file(func_to_be_decorated, *args, **kwargs): # Note that since we have used @decorator, *args, and **kwargs have # already been resolved to match the function signature of func. This # means default values have been propagated. For example, the function # func(x, y, a=1, b=2, **kwargs) if called as func(0,1,b=5,c=10) would # have args=(0,1,1,5) and kwargs={'c':10}. # First we parse the arguments of the decorator. The path_arg could # be an positional argument or a keyword argument. Even if it is try: # path_arg is a required positional argument # This works precisely because we are using @decorator path = args[path_arg] except TypeError: # path_arg is a keyword argument. It is "required" in the sense # that it must exist, according to the decorator specification, # It can exist in `kwargs` by a developer specified default value # or it could have been explicitly set by the user. try: path = kwargs[path_arg] except KeyError: # Could not find the keyword. Thus, no default was specified # in the function signature and the user did not provide it. msg = 'Missing required keyword argument: {0}' raise nx.NetworkXError(msg.format(path_arg)) else: is_kwarg = True except IndexError: # A "required" argument was missing. This can only happen if # the decorator of the function was incorrectly specified. # So this probably is not a user error, but a developer error. msg = "path_arg of open_file decorator is incorrect" raise nx.NetworkXError(msg) else: is_kwarg = False # Now we have the path_arg. There are two types of input to consider: # 1) string representing a path that should be opened # 2) an already opened file object if is_string_like(path): ext = splitext(path)[1] fobj = _dispatch_dict[ext](path, mode=mode) close_fobj = True elif hasattr(path, 'read'): # path is already a file-like object fobj = path close_fobj = False elif Path is not None and isinstance(path, Path): # path is a pathlib reference to a filename fobj = _dispatch_dict[path.suffix](str(path), mode=mode) close_fobj = True else: # could be None, in which case the algorithm will deal with it fobj = path close_fobj = False # Insert file object into args or kwargs. if is_kwarg: new_args = args kwargs[path_arg] = fobj else: # args is a tuple, so we must convert to list before modifying it. new_args = list(args) new_args[path_arg] = fobj # Finally, we call the original function, making sure to close the fobj try: result = func_to_be_decorated(*new_args, **kwargs) finally: if close_fobj: fobj.close() return result return _open_file
[文档]def nodes_or_number(which_args): """Decorator to allow number of nodes or container of nodes. Parameters ---------- which_args : int or sequence of ints Location of the node arguments in args. Even if the argument is a named positional argument (with a default value), you must specify its index as a positional argument. If more than one node argument is allowed, can be a list of locations. Returns ------- _nodes_or_numbers : function Function which replaces int args with ranges. Examples -------- Decorate functions like this:: @nodes_or_number(0) def empty_graph(nodes): pass @nodes_or_number([0,1]) def grid_2d_graph(m1, m2, periodic=False): pass @nodes_or_number(1) def full_rary_tree(r, n) # r is a number. n can be a number of a list of nodes pass """ @decorator def _nodes_or_number(func_to_be_decorated, *args, **kw): # form tuple of arg positions to be converted. try: iter_wa = iter(which_args) except TypeError: iter_wa = (which_args,) # change each argument in turn new_args = list(args) for i in iter_wa: n = args[i] try: nodes = list(range(n)) except TypeError: nodes = tuple(n) else: if n < 0: msg = "Negative number of nodes not valid: %i" % n raise nx.NetworkXError(msg) new_args[i] = (n, nodes) return func_to_be_decorated(*new_args, **kw) return _nodes_or_number
[文档]def preserve_random_state(func): """ Decorator to preserve the numpy.random state during a function. Parameters ---------- func : function function around which to preserve the random state. Returns ------- wrapper : function Function which wraps the input function by saving the state before calling the function and restoring the function afterward. Examples -------- Decorate functions like this:: @preserve_random_state def do_random_stuff(x, y): return x + y * numpy.random.random() Notes ----- If numpy.random is not importable, the state is not saved or restored. """ try: from numpy.random import get_state, seed, set_state @contextmanager def save_random_state(): state = get_state() try: yield finally: set_state(state) def wrapper(*args, **kwargs): with save_random_state(): seed(1234567890) return func(*args, **kwargs) wrapper.__name__ = func.__name__ return wrapper except ImportError: return func
[文档]def random_state(random_state_index): """Decorator to generate a numpy.random.RandomState instance. Argument position `random_state_index` is processed by create_random_state. The result is a numpy.random.RandomState instance. Parameters ---------- random_state_index : int Location of the random_state argument in args that is to be used to generate the numpy.random.RandomState instance. Even if the argument is a named positional argument (with a default value), you must specify its index as a positional argument. Returns ------- _random_state : function Function whose random_state keyword argument is a RandomState instance. Examples -------- Decorate functions like this:: @np_random_state(0) def random_float(random_state=None): return random_state.rand() @np_random_state(1) def random_array(dims, random_state=1): return random_state.rand(*dims) See Also -------- py_random_state """ @decorator def _random_state(func, *args, **kwargs): # Parse the decorator arguments. try: random_state_arg = args[random_state_index] except TypeError: raise nx.NetworkXError("random_state_index must be an integer") except IndexError: raise nx.NetworkXError("random_state_index is incorrect") # Create a numpy.random.RandomState instance random_state = create_random_state(random_state_arg) # args is a tuple, so we must convert to list before modifying it. new_args = list(args) new_args[random_state_index] = random_state return func(*new_args, **kwargs) return _random_state
np_random_state = random_state def py_random_state(random_state_index): """Decorator to generate a random.Random instance (or equiv). Argument position `random_state_index` processed by create_py_random_state. The result is either a random.Random instance, or numpy.random.RandomState instance with additional attributes to mimic basic methods of Random. Parameters ---------- random_state_index : int Location of the random_state argument in args that is to be used to generate the numpy.random.RandomState instance. Even if the argument is a named positional argument (with a default value), you must specify its index as a positional argument. Returns ------- _random_state : function Function whose random_state keyword argument is a RandomState instance. Examples -------- Decorate functions like this:: @py_random_state(0) def random_float(random_state=None): return random_state.rand() @py_random_state(1) def random_array(dims, random_state=1): return random_state.rand(*dims) See Also -------- np_random_state """ @decorator def _random_state(func, *args, **kwargs): # Parse the decorator arguments. try: random_state_arg = args[random_state_index] except TypeError: raise nx.NetworkXError("random_state_index must be an integer") except IndexError: raise nx.NetworkXError("random_state_index is incorrect") # Create a numpy.random.RandomState instance random_state = create_py_random_state(random_state_arg) # args is a tuple, so we must convert to list before modifying it. new_args = list(args) new_args[random_state_index] = random_state return func(*new_args, **kwargs) return _random_state