decorators.py 43 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209
  1. import bz2
  2. import collections
  3. import gzip
  4. import inspect
  5. import itertools
  6. import re
  7. from collections import defaultdict
  8. from contextlib import contextmanager
  9. from os.path import splitext
  10. from pathlib import Path
  11. import networkx as nx
  12. from networkx.utils import create_py_random_state, create_random_state
  13. __all__ = [
  14. "not_implemented_for",
  15. "open_file",
  16. "nodes_or_number",
  17. "np_random_state",
  18. "py_random_state",
  19. "argmap",
  20. ]
  21. def not_implemented_for(*graph_types):
  22. """Decorator to mark algorithms as not implemented
  23. Parameters
  24. ----------
  25. graph_types : container of strings
  26. Entries must be one of "directed", "undirected", "multigraph", or "graph".
  27. Returns
  28. -------
  29. _require : function
  30. The decorated function.
  31. Raises
  32. ------
  33. NetworkXNotImplemented
  34. If any of the packages cannot be imported
  35. Notes
  36. -----
  37. Multiple types are joined logically with "and".
  38. For "or" use multiple @not_implemented_for() lines.
  39. Examples
  40. --------
  41. Decorate functions like this::
  42. @not_implemented_for("directed")
  43. def sp_function(G):
  44. pass
  45. # rule out MultiDiGraph
  46. @not_implemented_for("directed","multigraph")
  47. def sp_np_function(G):
  48. pass
  49. # rule out all except DiGraph
  50. @not_implemented_for("undirected")
  51. @not_implemented_for("multigraph")
  52. def sp_np_function(G):
  53. pass
  54. """
  55. if ("directed" in graph_types) and ("undirected" in graph_types):
  56. raise ValueError("Function not implemented on directed AND undirected graphs?")
  57. if ("multigraph" in graph_types) and ("graph" in graph_types):
  58. raise ValueError("Function not implemented on graph AND multigraphs?")
  59. if not set(graph_types) < {"directed", "undirected", "multigraph", "graph"}:
  60. raise KeyError(
  61. "use one or more of directed, undirected, multigraph, graph. "
  62. f"You used {graph_types}"
  63. )
  64. # 3-way logic: True if "directed" input, False if "undirected" input, else None
  65. dval = ("directed" in graph_types) or "undirected" not in graph_types and None
  66. mval = ("multigraph" in graph_types) or "graph" not in graph_types and None
  67. errmsg = f"not implemented for {' '.join(graph_types)} type"
  68. def _not_implemented_for(g):
  69. if (mval is None or mval == g.is_multigraph()) and (
  70. dval is None or dval == g.is_directed()
  71. ):
  72. raise nx.NetworkXNotImplemented(errmsg)
  73. return g
  74. return argmap(_not_implemented_for, 0)
  75. # To handle new extensions, define a function accepting a `path` and `mode`.
  76. # Then add the extension to _dispatch_dict.
  77. fopeners = {
  78. ".gz": gzip.open,
  79. ".gzip": gzip.open,
  80. ".bz2": bz2.BZ2File,
  81. }
  82. _dispatch_dict = defaultdict(lambda: open, **fopeners)
  83. def open_file(path_arg, mode="r"):
  84. """Decorator to ensure clean opening and closing of files.
  85. Parameters
  86. ----------
  87. path_arg : string or int
  88. Name or index of the argument that is a path.
  89. mode : str
  90. String for opening mode.
  91. Returns
  92. -------
  93. _open_file : function
  94. Function which cleanly executes the io.
  95. Examples
  96. --------
  97. Decorate functions like this::
  98. @open_file(0,"r")
  99. def read_function(pathname):
  100. pass
  101. @open_file(1,"w")
  102. def write_function(G, pathname):
  103. pass
  104. @open_file(1,"w")
  105. def write_function(G, pathname="graph.dot"):
  106. pass
  107. @open_file("pathname","w")
  108. def write_function(G, pathname="graph.dot"):
  109. pass
  110. @open_file("path", "w+")
  111. def another_function(arg, **kwargs):
  112. path = kwargs["path"]
  113. pass
  114. Notes
  115. -----
  116. Note that this decorator solves the problem when a path argument is
  117. specified as a string, but it does not handle the situation when the
  118. function wants to accept a default of None (and then handle it).
  119. Here is an example of how to handle this case::
  120. @open_file("path")
  121. def some_function(arg1, arg2, path=None):
  122. if path is None:
  123. fobj = tempfile.NamedTemporaryFile(delete=False)
  124. else:
  125. # `path` could have been a string or file object or something
  126. # similar. In any event, the decorator has given us a file object
  127. # and it will close it for us, if it should.
  128. fobj = path
  129. try:
  130. fobj.write("blah")
  131. finally:
  132. if path is None:
  133. fobj.close()
  134. Normally, we'd want to use "with" to ensure that fobj gets closed.
  135. However, the decorator will make `path` a file object for us,
  136. and using "with" would undesirably close that file object.
  137. Instead, we use a try block, as shown above.
  138. When we exit the function, fobj will be closed, if it should be, by the decorator.
  139. """
  140. def _open_file(path):
  141. # Now we have the path_arg. There are two types of input to consider:
  142. # 1) string representing a path that should be opened
  143. # 2) an already opened file object
  144. if isinstance(path, str):
  145. ext = splitext(path)[1]
  146. elif isinstance(path, Path):
  147. # path is a pathlib reference to a filename
  148. ext = path.suffix
  149. path = str(path)
  150. else:
  151. # could be None, or a file handle, in which case the algorithm will deal with it
  152. return path, lambda: None
  153. fobj = _dispatch_dict[ext](path, mode=mode)
  154. return fobj, lambda: fobj.close()
  155. return argmap(_open_file, path_arg, try_finally=True)
  156. def nodes_or_number(which_args):
  157. """Decorator to allow number of nodes or container of nodes.
  158. With this decorator, the specified argument can be either a number or a container
  159. of nodes. If it is a number, the nodes used are `range(n)`.
  160. This allows `nx.complete_graph(50)` in place of `nx.complete_graph(list(range(50)))`.
  161. And it also allows `nx.complete_graph(any_list_of_nodes)`.
  162. Parameters
  163. ----------
  164. which_args : string or int or sequence of strings or ints
  165. If string, the name of the argument to be treated.
  166. If int, the index of the argument to be treated.
  167. If more than one node argument is allowed, can be a list of locations.
  168. Returns
  169. -------
  170. _nodes_or_numbers : function
  171. Function which replaces int args with ranges.
  172. Examples
  173. --------
  174. Decorate functions like this::
  175. @nodes_or_number("nodes")
  176. def empty_graph(nodes):
  177. # nodes is converted to a list of nodes
  178. @nodes_or_number(0)
  179. def empty_graph(nodes):
  180. # nodes is converted to a list of nodes
  181. @nodes_or_number(["m1", "m2"])
  182. def grid_2d_graph(m1, m2, periodic=False):
  183. # m1 and m2 are each converted to a list of nodes
  184. @nodes_or_number([0, 1])
  185. def grid_2d_graph(m1, m2, periodic=False):
  186. # m1 and m2 are each converted to a list of nodes
  187. @nodes_or_number(1)
  188. def full_rary_tree(r, n)
  189. # presumably r is a number. It is not handled by this decorator.
  190. # n is converted to a list of nodes
  191. """
  192. def _nodes_or_number(n):
  193. try:
  194. nodes = list(range(n))
  195. except TypeError:
  196. nodes = tuple(n)
  197. else:
  198. if n < 0:
  199. raise nx.NetworkXError(f"Negative number of nodes not valid: {n}")
  200. return (n, nodes)
  201. try:
  202. iter_wa = iter(which_args)
  203. except TypeError:
  204. iter_wa = (which_args,)
  205. return argmap(_nodes_or_number, *iter_wa)
  206. def np_random_state(random_state_argument):
  207. """Decorator to generate a `numpy.random.RandomState` instance.
  208. The decorator processes the argument indicated by `random_state_argument`
  209. using :func:`nx.utils.create_random_state`.
  210. The argument value can be a seed (integer), or a `numpy.random.RandomState`
  211. instance or (`None` or `numpy.random`). The latter options use the glocal
  212. random number generator used by `numpy.random`.
  213. The result is a `numpy.random.RandomState` instance.
  214. Parameters
  215. ----------
  216. random_state_argument : string or int
  217. The name or index of the argument to be converted
  218. to a `numpy.random.RandomState` instance.
  219. Returns
  220. -------
  221. _random_state : function
  222. Function whose random_state keyword argument is a RandomState instance.
  223. Examples
  224. --------
  225. Decorate functions like this::
  226. @np_random_state("seed")
  227. def random_float(seed=None):
  228. return seed.rand()
  229. @np_random_state(0)
  230. def random_float(rng=None):
  231. return rng.rand()
  232. @np_random_state(1)
  233. def random_array(dims, random_state=1):
  234. return random_state.rand(*dims)
  235. See Also
  236. --------
  237. py_random_state
  238. """
  239. return argmap(create_random_state, random_state_argument)
  240. def py_random_state(random_state_argument):
  241. """Decorator to generate a random.Random instance (or equiv).
  242. The decorator processes the argument indicated by `random_state_argument`
  243. using :func:`nx.utils.create_py_random_state`.
  244. The argument value can be a seed (integer), or a random number generator::
  245. If int, return a random.Random instance set with seed=int.
  246. If random.Random instance, return it.
  247. If None or the `random` package, return the global random number
  248. generator used by `random`.
  249. If np.random package, return the global numpy random number
  250. generator wrapped in a PythonRandomInterface class.
  251. If np.random.RandomState instance, return it wrapped in
  252. PythonRandomInterface
  253. If a PythonRandomInterface instance, return it
  254. Parameters
  255. ----------
  256. random_state_argument : string or int
  257. The name of the argument or the index of the argument in args that is
  258. to be converted to the random.Random instance or numpy.random.RandomState
  259. instance that mimics basic methods of random.Random.
  260. Returns
  261. -------
  262. _random_state : function
  263. Function whose random_state_argument is converted to a Random instance.
  264. Examples
  265. --------
  266. Decorate functions like this::
  267. @py_random_state("random_state")
  268. def random_float(random_state=None):
  269. return random_state.rand()
  270. @py_random_state(0)
  271. def random_float(rng=None):
  272. return rng.rand()
  273. @py_random_state(1)
  274. def random_array(dims, seed=12345):
  275. return seed.rand(*dims)
  276. See Also
  277. --------
  278. np_random_state
  279. """
  280. return argmap(create_py_random_state, random_state_argument)
  281. class argmap:
  282. """A decorator to apply a map to arguments before calling the function
  283. This class provides a decorator that maps (transforms) arguments of the function
  284. before the function is called. Thus for example, we have similar code
  285. in many functions to determine whether an argument is the number of nodes
  286. to be created, or a list of nodes to be handled. The decorator provides
  287. the code to accept either -- transforming the indicated argument into a
  288. list of nodes before the actual function is called.
  289. This decorator class allows us to process single or multiple arguments.
  290. The arguments to be processed can be specified by string, naming the argument,
  291. or by index, specifying the item in the args list.
  292. Parameters
  293. ----------
  294. func : callable
  295. The function to apply to arguments
  296. *args : iterable of (int, str or tuple)
  297. A list of parameters, specified either as strings (their names), ints
  298. (numerical indices) or tuples, which may contain ints, strings, and
  299. (recursively) tuples. Each indicates which parameters the decorator
  300. should map. Tuples indicate that the map function takes (and returns)
  301. multiple parameters in the same order and nested structure as indicated
  302. here.
  303. try_finally : bool (default: False)
  304. When True, wrap the function call in a try-finally block with code
  305. for the finally block created by `func`. This is used when the map
  306. function constructs an object (like a file handle) that requires
  307. post-processing (like closing).
  308. Note: try_finally decorators cannot be used to decorate generator
  309. functions.
  310. Examples
  311. --------
  312. Most of these examples use `@argmap(...)` to apply the decorator to
  313. the function defined on the next line.
  314. In the NetworkX codebase however, `argmap` is used within a function to
  315. construct a decorator. That is, the decorator defines a mapping function
  316. and then uses `argmap` to build and return a decorated function.
  317. A simple example is a decorator that specifies which currency to report money.
  318. The decorator (named `convert_to`) would be used like::
  319. @convert_to("US_Dollars", "income")
  320. def show_me_the_money(name, income):
  321. print(f"{name} : {income}")
  322. And the code to create the decorator might be::
  323. def convert_to(currency, which_arg):
  324. def _convert(amount):
  325. if amount.currency != currency:
  326. amount = amount.to_currency(currency)
  327. return amount
  328. return argmap(_convert, which_arg)
  329. Despite this common idiom for argmap, most of the following examples
  330. use the `@argmap(...)` idiom to save space.
  331. Here's an example use of argmap to sum the elements of two of the functions
  332. arguments. The decorated function::
  333. @argmap(sum, "xlist", "zlist")
  334. def foo(xlist, y, zlist):
  335. return xlist - y + zlist
  336. is syntactic sugar for::
  337. def foo(xlist, y, zlist):
  338. x = sum(xlist)
  339. z = sum(zlist)
  340. return x - y + z
  341. and is equivalent to (using argument indexes)::
  342. @argmap(sum, "xlist", 2)
  343. def foo(xlist, y, zlist):
  344. return xlist - y + zlist
  345. or::
  346. @argmap(sum, "zlist", 0)
  347. def foo(xlist, y, zlist):
  348. return xlist - y + zlist
  349. Transforming functions can be applied to multiple arguments, such as::
  350. def swap(x, y):
  351. return y, x
  352. # the 2-tuple tells argmap that the map `swap` has 2 inputs/outputs.
  353. @argmap(swap, ("a", "b")):
  354. def foo(a, b, c):
  355. return a / b * c
  356. is equivalent to::
  357. def foo(a, b, c):
  358. a, b = swap(a, b)
  359. return a / b * c
  360. More generally, the applied arguments can be nested tuples of strings or ints.
  361. The syntax `@argmap(some_func, ("a", ("b", "c")))` would expect `some_func` to
  362. accept 2 inputs with the second expected to be a 2-tuple. It should then return
  363. 2 outputs with the second a 2-tuple. The returns values would replace input "a"
  364. "b" and "c" respectively. Similarly for `@argmap(some_func, (0, ("b", 2)))`.
  365. Also, note that an index larger than the number of named parameters is allowed
  366. for variadic functions. For example::
  367. def double(a):
  368. return 2 * a
  369. @argmap(double, 3)
  370. def overflow(a, *args):
  371. return a, args
  372. print(overflow(1, 2, 3, 4, 5, 6)) # output is 1, (2, 3, 8, 5, 6)
  373. **Try Finally**
  374. Additionally, this `argmap` class can be used to create a decorator that
  375. initiates a try...finally block. The decorator must be written to return
  376. both the transformed argument and a closing function.
  377. This feature was included to enable the `open_file` decorator which might
  378. need to close the file or not depending on whether it had to open that file.
  379. This feature uses the keyword-only `try_finally` argument to `@argmap`.
  380. For example this map opens a file and then makes sure it is closed::
  381. def open_file(fn):
  382. f = open(fn)
  383. return f, lambda: f.close()
  384. The decorator applies that to the function `foo`::
  385. @argmap(open_file, "file", try_finally=True)
  386. def foo(file):
  387. print(file.read())
  388. is syntactic sugar for::
  389. def foo(file):
  390. file, close_file = open_file(file)
  391. try:
  392. print(file.read())
  393. finally:
  394. close_file()
  395. and is equivalent to (using indexes)::
  396. @argmap(open_file, 0, try_finally=True)
  397. def foo(file):
  398. print(file.read())
  399. Here's an example of the try_finally feature used to create a decorator::
  400. def my_closing_decorator(which_arg):
  401. def _opener(path):
  402. if path is None:
  403. path = open(path)
  404. fclose = path.close
  405. else:
  406. # assume `path` handles the closing
  407. fclose = lambda: None
  408. return path, fclose
  409. return argmap(_opener, which_arg, try_finally=True)
  410. which can then be used as::
  411. @my_closing_decorator("file")
  412. def fancy_reader(file=None):
  413. # this code doesn't need to worry about closing the file
  414. print(file.read())
  415. Decorators with try_finally = True cannot be used with generator functions,
  416. because the `finally` block is evaluated before the generator is exhausted::
  417. @argmap(open_file, "file", try_finally=True)
  418. def file_to_lines(file):
  419. for line in file.readlines():
  420. yield line
  421. is equivalent to::
  422. def file_to_lines_wrapped(file):
  423. for line in file.readlines():
  424. yield line
  425. def file_to_lines_wrapper(file):
  426. try:
  427. file = open_file(file)
  428. return file_to_lines_wrapped(file)
  429. finally:
  430. file.close()
  431. which behaves similarly to::
  432. def file_to_lines_whoops(file):
  433. file = open_file(file)
  434. file.close()
  435. for line in file.readlines():
  436. yield line
  437. because the `finally` block of `file_to_lines_wrapper` is executed before
  438. the caller has a chance to exhaust the iterator.
  439. Notes
  440. -----
  441. An object of this class is callable and intended to be used when
  442. defining a decorator. Generally, a decorator takes a function as input
  443. and constructs a function as output. Specifically, an `argmap` object
  444. returns the input function decorated/wrapped so that specified arguments
  445. are mapped (transformed) to new values before the decorated function is called.
  446. As an overview, the argmap object returns a new function with all the
  447. dunder values of the original function (like `__doc__`, `__name__`, etc).
  448. Code for this decorated function is built based on the original function's
  449. signature. It starts by mapping the input arguments to potentially new
  450. values. Then it calls the decorated function with these new values in place
  451. of the indicated arguments that have been mapped. The return value of the
  452. original function is then returned. This new function is the function that
  453. is actually called by the user.
  454. Three additional features are provided.
  455. 1) The code is lazily compiled. That is, the new function is returned
  456. as an object without the code compiled, but with all information
  457. needed so it can be compiled upon it's first invocation. This saves
  458. time on import at the cost of additional time on the first call of
  459. the function. Subsequent calls are then just as fast as normal.
  460. 2) If the "try_finally" keyword-only argument is True, a try block
  461. follows each mapped argument, matched on the other side of the wrapped
  462. call, by a finally block closing that mapping. We expect func to return
  463. a 2-tuple: the mapped value and a function to be called in the finally
  464. clause. This feature was included so the `open_file` decorator could
  465. provide a file handle to the decorated function and close the file handle
  466. after the function call. It even keeps track of whether to close the file
  467. handle or not based on whether it had to open the file or the input was
  468. already open. So, the decorated function does not need to include any
  469. code to open or close files.
  470. 3) The maps applied can process multiple arguments. For example,
  471. you could swap two arguments using a mapping, or transform
  472. them to their sum and their difference. This was included to allow
  473. a decorator in the `quality.py` module that checks that an input
  474. `partition` is a valid partition of the nodes of the input graph `G`.
  475. In this example, the map has inputs `(G, partition)`. After checking
  476. for a valid partition, the map either raises an exception or leaves
  477. the inputs unchanged. Thus many functions that make this check can
  478. use the decorator rather than copy the checking code into each function.
  479. More complicated nested argument structures are described below.
  480. The remaining notes describe the code structure and methods for this
  481. class in broad terms to aid in understanding how to use it.
  482. Instantiating an `argmap` object simply stores the mapping function and
  483. the input identifiers of which arguments to map. The resulting decorator
  484. is ready to use this map to decorate any function. Calling that object
  485. (`argmap.__call__`, but usually done via `@my_decorator`) a lazily
  486. compiled thin wrapper of the decorated function is constructed,
  487. wrapped with the necessary function dunder attributes like `__doc__`
  488. and `__name__`. That thinly wrapped function is returned as the
  489. decorated function. When that decorated function is called, the thin
  490. wrapper of code calls `argmap._lazy_compile` which compiles the decorated
  491. function (using `argmap.compile`) and replaces the code of the thin
  492. wrapper with the newly compiled code. This saves the compilation step
  493. every import of networkx, at the cost of compiling upon the first call
  494. to the decorated function.
  495. When the decorated function is compiled, the code is recursively assembled
  496. using the `argmap.assemble` method. The recursive nature is needed in
  497. case of nested decorators. The result of the assembly is a number of
  498. useful objects.
  499. sig : the function signature of the original decorated function as
  500. constructed by :func:`argmap.signature`. This is constructed
  501. using `inspect.signature` but enhanced with attribute
  502. strings `sig_def` and `sig_call`, and other information
  503. specific to mapping arguments of this function.
  504. This information is used to construct a string of code defining
  505. the new decorated function.
  506. wrapped_name : a unique internally used name constructed by argmap
  507. for the decorated function.
  508. functions : a dict of the functions used inside the code of this
  509. decorated function, to be used as `globals` in `exec`.
  510. This dict is recursively updated to allow for nested decorating.
  511. mapblock : code (as a list of strings) to map the incoming argument
  512. values to their mapped values.
  513. finallys : code (as a list of strings) to provide the possibly nested
  514. set of finally clauses if needed.
  515. mutable_args : a bool indicating whether the `sig.args` tuple should be
  516. converted to a list so mutation can occur.
  517. After this recursive assembly process, the `argmap.compile` method
  518. constructs code (as strings) to convert the tuple `sig.args` to a list
  519. if needed. It joins the defining code with appropriate indents and
  520. compiles the result. Finally, this code is evaluated and the original
  521. wrapper's implementation is replaced with the compiled version (see
  522. `argmap._lazy_compile` for more details).
  523. Other `argmap` methods include `_name` and `_count` which allow internally
  524. generated names to be unique within a python session.
  525. The methods `_flatten` and `_indent` process the nested lists of strings
  526. into properly indented python code ready to be compiled.
  527. More complicated nested tuples of arguments also allowed though
  528. usually not used. For the simple 2 argument case, the argmap
  529. input ("a", "b") implies the mapping function will take 2 arguments
  530. and return a 2-tuple of mapped values. A more complicated example
  531. with argmap input `("a", ("b", "c"))` requires the mapping function
  532. take 2 inputs, with the second being a 2-tuple. It then must output
  533. the 3 mapped values in the same nested structure `(newa, (newb, newc))`.
  534. This level of generality is not often needed, but was convenient
  535. to implement when handling the multiple arguments.
  536. See Also
  537. --------
  538. not_implemented_for
  539. open_file
  540. nodes_or_number
  541. random_state
  542. py_random_state
  543. networkx.community.quality.require_partition
  544. require_partition
  545. """
  546. def __init__(self, func, *args, try_finally=False):
  547. self._func = func
  548. self._args = args
  549. self._finally = try_finally
  550. @staticmethod
  551. def _lazy_compile(func):
  552. """Compile the source of a wrapped function
  553. Assemble and compile the decorated function, and intrusively replace its
  554. code with the compiled version's. The thinly wrapped function becomes
  555. the decorated function.
  556. Parameters
  557. ----------
  558. func : callable
  559. A function returned by argmap.__call__ which is in the process
  560. of being called for the first time.
  561. Returns
  562. -------
  563. func : callable
  564. The same function, with a new __code__ object.
  565. Notes
  566. -----
  567. It was observed in NetworkX issue #4732 [1] that the import time of
  568. NetworkX was significantly bloated by the use of decorators: over half
  569. of the import time was being spent decorating functions. This was
  570. somewhat improved by a change made to the `decorator` library, at the
  571. cost of a relatively heavy-weight call to `inspect.Signature.bind`
  572. for each call to the decorated function.
  573. The workaround we arrived at is to do minimal work at the time of
  574. decoration. When the decorated function is called for the first time,
  575. we compile a function with the same function signature as the wrapped
  576. function. The resulting decorated function is faster than one made by
  577. the `decorator` library, so that the overhead of the first call is
  578. 'paid off' after a small number of calls.
  579. References
  580. ----------
  581. [1] https://github.com/networkx/networkx/issues/4732
  582. """
  583. real_func = func.__argmap__.compile(func.__wrapped__)
  584. func.__code__ = real_func.__code__
  585. func.__globals__.update(real_func.__globals__)
  586. func.__dict__.update(real_func.__dict__)
  587. return func
  588. def __call__(self, f):
  589. """Construct a lazily decorated wrapper of f.
  590. The decorated function will be compiled when it is called for the first time,
  591. and it will replace its own __code__ object so subsequent calls are fast.
  592. Parameters
  593. ----------
  594. f : callable
  595. A function to be decorated.
  596. Returns
  597. -------
  598. func : callable
  599. The decorated function.
  600. See Also
  601. --------
  602. argmap._lazy_compile
  603. """
  604. def func(*args, __wrapper=None, **kwargs):
  605. return argmap._lazy_compile(__wrapper)(*args, **kwargs)
  606. # standard function-wrapping stuff
  607. func.__name__ = f.__name__
  608. func.__doc__ = f.__doc__
  609. func.__defaults__ = f.__defaults__
  610. func.__kwdefaults__.update(f.__kwdefaults__ or {})
  611. func.__module__ = f.__module__
  612. func.__qualname__ = f.__qualname__
  613. func.__dict__.update(f.__dict__)
  614. func.__wrapped__ = f
  615. # now that we've wrapped f, we may have picked up some __dict__ or
  616. # __kwdefaults__ items that were set by a previous argmap. Thus, we set
  617. # these values after those update() calls.
  618. # If we attempt to access func from within itself, that happens through
  619. # a closure -- which trips an error when we replace func.__code__. The
  620. # standard workaround for functions which can't see themselves is to use
  621. # a Y-combinator, as we do here.
  622. func.__kwdefaults__["_argmap__wrapper"] = func
  623. # this self-reference is here because functools.wraps preserves
  624. # everything in __dict__, and we don't want to mistake a non-argmap
  625. # wrapper for an argmap wrapper
  626. func.__self__ = func
  627. # this is used to variously call self.assemble and self.compile
  628. func.__argmap__ = self
  629. if hasattr(f, "__argmap__"):
  630. func.__is_generator = f.__is_generator
  631. else:
  632. func.__is_generator = inspect.isgeneratorfunction(f)
  633. if self._finally and func.__is_generator:
  634. raise nx.NetworkXError("argmap cannot decorate generators with try_finally")
  635. return func
  636. __count = 0
  637. @classmethod
  638. def _count(cls):
  639. """Maintain a globally-unique identifier for function names and "file" names
  640. Note that this counter is a class method reporting a class variable
  641. so the count is unique within a Python session. It could differ from
  642. session to session for a specific decorator depending on the order
  643. that the decorators are created. But that doesn't disrupt `argmap`.
  644. This is used in two places: to construct unique variable names
  645. in the `_name` method and to construct unique fictitious filenames
  646. in the `_compile` method.
  647. Returns
  648. -------
  649. count : int
  650. An integer unique to this Python session (simply counts from zero)
  651. """
  652. cls.__count += 1
  653. return cls.__count
  654. _bad_chars = re.compile("[^a-zA-Z0-9_]")
  655. @classmethod
  656. def _name(cls, f):
  657. """Mangle the name of a function to be unique but somewhat human-readable
  658. The names are unique within a Python session and set using `_count`.
  659. Parameters
  660. ----------
  661. f : str or object
  662. Returns
  663. -------
  664. name : str
  665. The mangled version of `f.__name__` (if `f.__name__` exists) or `f`
  666. """
  667. f = f.__name__ if hasattr(f, "__name__") else f
  668. fname = re.sub(cls._bad_chars, "_", f)
  669. return f"argmap_{fname}_{cls._count()}"
  670. def compile(self, f):
  671. """Compile the decorated function.
  672. Called once for a given decorated function -- collects the code from all
  673. argmap decorators in the stack, and compiles the decorated function.
  674. Much of the work done here uses the `assemble` method to allow recursive
  675. treatment of multiple argmap decorators on a single decorated function.
  676. That flattens the argmap decorators, collects the source code to construct
  677. a single decorated function, then compiles/executes/returns that function.
  678. The source code for the decorated function is stored as an attribute
  679. `_code` on the function object itself.
  680. Note that Python's `compile` function requires a filename, but this
  681. code is constructed without a file, so a fictitious filename is used
  682. to describe where the function comes from. The name is something like:
  683. "argmap compilation 4".
  684. Parameters
  685. ----------
  686. f : callable
  687. The function to be decorated
  688. Returns
  689. -------
  690. func : callable
  691. The decorated file
  692. """
  693. sig, wrapped_name, functions, mapblock, finallys, mutable_args = self.assemble(
  694. f
  695. )
  696. call = f"{sig.call_sig.format(wrapped_name)}#"
  697. mut_args = f"{sig.args} = list({sig.args})" if mutable_args else ""
  698. body = argmap._indent(sig.def_sig, mut_args, mapblock, call, finallys)
  699. code = "\n".join(body)
  700. locl = {}
  701. globl = dict(functions.values())
  702. filename = f"{self.__class__} compilation {self._count()}"
  703. compiled = compile(code, filename, "exec")
  704. exec(compiled, globl, locl)
  705. func = locl[sig.name]
  706. func._code = code
  707. return func
  708. def assemble(self, f):
  709. """Collects components of the source for the decorated function wrapping f.
  710. If `f` has multiple argmap decorators, we recursively assemble the stack of
  711. decorators into a single flattened function.
  712. This method is part of the `compile` method's process yet separated
  713. from that method to allow recursive processing. The outputs are
  714. strings, dictionaries and lists that collect needed info to
  715. flatten any nested argmap-decoration.
  716. Parameters
  717. ----------
  718. f : callable
  719. The function to be decorated. If f is argmapped, we assemble it.
  720. Returns
  721. -------
  722. sig : argmap.Signature
  723. The function signature as an `argmap.Signature` object.
  724. wrapped_name : str
  725. The mangled name used to represent the wrapped function in the code
  726. being assembled.
  727. functions : dict
  728. A dictionary mapping id(g) -> (mangled_name(g), g) for functions g
  729. referred to in the code being assembled. These need to be present
  730. in the ``globals`` scope of ``exec`` when defining the decorated
  731. function.
  732. mapblock : list of lists and/or strings
  733. Code that implements mapping of parameters including any try blocks
  734. if needed. This code will precede the decorated function call.
  735. finallys : list of lists and/or strings
  736. Code that implements the finally blocks to post-process the
  737. arguments (usually close any files if needed) after the
  738. decorated function is called.
  739. mutable_args : bool
  740. True if the decorator needs to modify positional arguments
  741. via their indices. The compile method then turns the argument
  742. tuple into a list so that the arguments can be modified.
  743. """
  744. # first, we check if f is already argmapped -- if that's the case,
  745. # build up the function recursively.
  746. # > mapblock is generally a list of function calls of the sort
  747. # arg = func(arg)
  748. # in addition to some try-blocks if needed.
  749. # > finallys is a recursive list of finally blocks of the sort
  750. # finally:
  751. # close_func_1()
  752. # finally:
  753. # close_func_2()
  754. # > functions is a dict of functions used in the scope of our decorated
  755. # function. It will be used to construct globals used in compilation.
  756. # We make functions[id(f)] = name_of_f, f to ensure that a given
  757. # function is stored and named exactly once even if called by
  758. # nested decorators.
  759. if hasattr(f, "__argmap__") and f.__self__ is f:
  760. (
  761. sig,
  762. wrapped_name,
  763. functions,
  764. mapblock,
  765. finallys,
  766. mutable_args,
  767. ) = f.__argmap__.assemble(f.__wrapped__)
  768. functions = dict(functions) # shallow-copy just in case
  769. else:
  770. sig = self.signature(f)
  771. wrapped_name = self._name(f)
  772. mapblock, finallys = [], []
  773. functions = {id(f): (wrapped_name, f)}
  774. mutable_args = False
  775. if id(self._func) in functions:
  776. fname, _ = functions[id(self._func)]
  777. else:
  778. fname, _ = functions[id(self._func)] = self._name(self._func), self._func
  779. # this is a bit complicated -- we can call functions with a variety of
  780. # nested arguments, so long as their input and output are tuples with
  781. # the same nested structure. e.g. ("a", "b") maps arguments a and b.
  782. # A more complicated nesting like (0, (3, 4)) maps arguments 0, 3, 4
  783. # expecting the mapping to output new values in the same nested shape.
  784. # The ability to argmap multiple arguments was necessary for
  785. # the decorator `nx.algorithms.community.quality.require_partition`, and
  786. # while we're not taking full advantage of the ability to handle
  787. # multiply-nested tuples, it was convenient to implement this in
  788. # generality because the recursive call to `get_name` is necessary in
  789. # any case.
  790. applied = set()
  791. def get_name(arg, first=True):
  792. nonlocal mutable_args
  793. if isinstance(arg, tuple):
  794. name = ", ".join(get_name(x, False) for x in arg)
  795. return name if first else f"({name})"
  796. if arg in applied:
  797. raise nx.NetworkXError(f"argument {arg} is specified multiple times")
  798. applied.add(arg)
  799. if arg in sig.names:
  800. return sig.names[arg]
  801. elif isinstance(arg, str):
  802. if sig.kwargs is None:
  803. raise nx.NetworkXError(
  804. f"name {arg} is not a named parameter and this function doesn't have kwargs"
  805. )
  806. return f"{sig.kwargs}[{arg!r}]"
  807. else:
  808. if sig.args is None:
  809. raise nx.NetworkXError(
  810. f"index {arg} not a parameter index and this function doesn't have args"
  811. )
  812. mutable_args = True
  813. return f"{sig.args}[{arg - sig.n_positional}]"
  814. if self._finally:
  815. # here's where we handle try_finally decorators. Such a decorator
  816. # returns a mapped argument and a function to be called in a
  817. # finally block. This feature was required by the open_file
  818. # decorator. The below generates the code
  819. #
  820. # name, final = func(name) #<--append to mapblock
  821. # try: #<--append to mapblock
  822. # ... more argmapping and try blocks
  823. # return WRAPPED_FUNCTION(...)
  824. # ... more finally blocks
  825. # finally: #<--prepend to finallys
  826. # final() #<--prepend to finallys
  827. #
  828. for a in self._args:
  829. name = get_name(a)
  830. final = self._name(name)
  831. mapblock.append(f"{name}, {final} = {fname}({name})")
  832. mapblock.append("try:")
  833. finallys = ["finally:", f"{final}()#", "#", finallys]
  834. else:
  835. mapblock.extend(
  836. f"{name} = {fname}({name})" for name in map(get_name, self._args)
  837. )
  838. return sig, wrapped_name, functions, mapblock, finallys, mutable_args
  839. @classmethod
  840. def signature(cls, f):
  841. r"""Construct a Signature object describing `f`
  842. Compute a Signature so that we can write a function wrapping f with
  843. the same signature and call-type.
  844. Parameters
  845. ----------
  846. f : callable
  847. A function to be decorated
  848. Returns
  849. -------
  850. sig : argmap.Signature
  851. The Signature of f
  852. Notes
  853. -----
  854. The Signature is a namedtuple with names:
  855. name : a unique version of the name of the decorated function
  856. signature : the inspect.signature of the decorated function
  857. def_sig : a string used as code to define the new function
  858. call_sig : a string used as code to call the decorated function
  859. names : a dict keyed by argument name and index to the argument's name
  860. n_positional : the number of positional arguments in the signature
  861. args : the name of the VAR_POSITIONAL argument if any, i.e. \*theseargs
  862. kwargs : the name of the VAR_KEYWORDS argument if any, i.e. \*\*kwargs
  863. These named attributes of the signature are used in `assemble` and `compile`
  864. to construct a string of source code for the decorated function.
  865. """
  866. sig = inspect.signature(f, follow_wrapped=False)
  867. def_sig = []
  868. call_sig = []
  869. names = {}
  870. kind = None
  871. args = None
  872. kwargs = None
  873. npos = 0
  874. for i, param in enumerate(sig.parameters.values()):
  875. # parameters can be position-only, keyword-or-position, keyword-only
  876. # in any combination, but only in the order as above. we do edge
  877. # detection to add the appropriate punctuation
  878. prev = kind
  879. kind = param.kind
  880. if prev == param.POSITIONAL_ONLY != kind:
  881. # the last token was position-only, but this one isn't
  882. def_sig.append("/")
  883. if prev != param.KEYWORD_ONLY == kind != param.VAR_POSITIONAL:
  884. # param is the first keyword-only arg and isn't starred
  885. def_sig.append("*")
  886. # star arguments as appropriate
  887. if kind == param.VAR_POSITIONAL:
  888. name = "*" + param.name
  889. args = param.name
  890. count = 0
  891. elif kind == param.VAR_KEYWORD:
  892. name = "**" + param.name
  893. kwargs = param.name
  894. count = 0
  895. else:
  896. names[i] = names[param.name] = param.name
  897. name = param.name
  898. count = 1
  899. # assign to keyword-only args in the function call
  900. if kind == param.KEYWORD_ONLY:
  901. call_sig.append(f"{name} = {name}")
  902. else:
  903. npos += count
  904. call_sig.append(name)
  905. def_sig.append(name)
  906. fname = cls._name(f)
  907. def_sig = f'def {fname}({", ".join(def_sig)}):'
  908. call_sig = f"return {{}}({', '.join(call_sig)})"
  909. return cls.Signature(fname, sig, def_sig, call_sig, names, npos, args, kwargs)
  910. Signature = collections.namedtuple(
  911. "Signature",
  912. [
  913. "name",
  914. "signature",
  915. "def_sig",
  916. "call_sig",
  917. "names",
  918. "n_positional",
  919. "args",
  920. "kwargs",
  921. ],
  922. )
  923. @staticmethod
  924. def _flatten(nestlist, visited):
  925. """flattens a recursive list of lists that doesn't have cyclic references
  926. Parameters
  927. ----------
  928. nestlist : iterable
  929. A recursive list of objects to be flattened into a single iterable
  930. visited : set
  931. A set of object ids which have been walked -- initialize with an
  932. empty set
  933. Yields
  934. ------
  935. Non-list objects contained in nestlist
  936. """
  937. for thing in nestlist:
  938. if isinstance(thing, list):
  939. if id(thing) in visited:
  940. raise ValueError("A cycle was found in nestlist. Be a tree.")
  941. else:
  942. visited.add(id(thing))
  943. yield from argmap._flatten(thing, visited)
  944. else:
  945. yield thing
  946. _tabs = " " * 64
  947. @staticmethod
  948. def _indent(*lines):
  949. """Indent list of code lines to make executable Python code
  950. Indents a tree-recursive list of strings, following the rule that one
  951. space is added to the tab after a line that ends in a colon, and one is
  952. removed after a line that ends in an hashmark.
  953. Parameters
  954. ----------
  955. *lines : lists and/or strings
  956. A recursive list of strings to be assembled into properly indented
  957. code.
  958. Returns
  959. -------
  960. code : str
  961. Examples
  962. --------
  963. argmap._indent(*["try:", "try:", "pass#", "finally:", "pass#", "#",
  964. "finally:", "pass#"])
  965. renders to
  966. '''try:
  967. try:
  968. pass#
  969. finally:
  970. pass#
  971. #
  972. finally:
  973. pass#'''
  974. """
  975. depth = 0
  976. for line in argmap._flatten(lines, set()):
  977. yield f"{argmap._tabs[:depth]}{line}"
  978. depth += (line[-1:] == ":") - (line[-1:] == "#")