diff --git a/openupgradelib/openupgrade.py b/openupgradelib/openupgrade.py index daf0ca27..8d57e6da 100644 --- a/openupgradelib/openupgrade.py +++ b/openupgradelib/openupgrade.py @@ -7,6 +7,7 @@ import inspect import uuid import logging as _logging_module +import functools from datetime import datetime try: from StringIO import StringIO @@ -129,6 +130,7 @@ def do_raise(error): __all__ = [ 'migrate', 'logging', + 'progress', 'load_data', 'add_fields', 'copy_columns', @@ -1876,6 +1878,52 @@ def wrapped_function(cr, version): return wrap +def progress(index=1, prefix=None): + """This is a decorator for any sub function called in an OpenUpgrade script + (pre, post or end migration script). + + Decorate functions that may take time. If a function is decorated, we + provide an iterable argument that will be looped on and call the original + function, while a progress bar will be displayed. + + Function: 28% (152 of 529) I##### I Elapsed Time: 0:00:03 ETA: 0:01:32 + + :param index: Index of the argument to be used as iterable. Default to the + second argument. It will pass each of the elements of the iterable in the + same place. + :param title: Optional title for prefixing the progress bar. If not + specified, the function name will be used. + + Typical use:: + + @openupgrade.progress() + def migrate_some_stuff(env, record) + # some custom code + ... + + @openupgrade.migrate() + def migrate(env, version): + records = ... # get an iterable + migrate_some_stuff(env, records) + """ + def wrap(func): + @functools.wraps(func) + def wrapped_function(*args, **kwargs): + import progressbar + elems = args[index] + prefix2 = prefix or str(func.__name__) + ": " + with progressbar.ProgressBar( + prefix=prefix2, max_value=len(elems) + ) as bar: + for elem in elems: + args_copy = list(args) + args_copy[index] = elem + func(*args_copy, **kwargs) + bar.update(bar.value + 1) + return wrapped_function + return wrap + + def move_field_m2o( cr, pool, registry_old_model, field_old_model, m2o_field_old_model, diff --git a/requirements.txt b/requirements.txt index df6b4ce6..41a54a9b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,5 @@ coveralls flake8 pep8-naming lxml<=4.3.4 +progressbar2 psycopg2==2.7.3.1 diff --git a/setup.py b/setup.py index 3e6894b0..06f58c90 100755 --- a/setup.py +++ b/setup.py @@ -36,7 +36,7 @@ packages=['openupgradelib'], package_dir={'openupgradelib': 'openupgradelib'}, include_package_data=True, - install_requires=["lxml", "cssselect"], + install_requires=["lxml", "cssselect", "progressbar2"], license=openupgradelib.__license__, zip_safe=False, keywords='openupgradelib',