From 9ba7e7a7063d077f3a0c7b2715882850d8afdccc Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 13 Oct 2020 03:49:14 -0600 Subject: [PATCH 1/3] Option to collapse iterable AST types --- pragma/collapse_literals.py | 9 ++++++++- pragma/core/resolve/__init__.py | 2 ++ pragma/core/transformer.py | 5 ++++- tests/pytest.ini | 4 ++++ tests/test_collapse_literals.py | 14 ++++++++++++++ 5 files changed, 32 insertions(+), 2 deletions(-) create mode 100644 tests/pytest.ini diff --git a/pragma/collapse_literals.py b/pragma/collapse_literals.py index 51968b3..1bea702 100644 --- a/pragma/collapse_literals.py +++ b/pragma/collapse_literals.py @@ -1,17 +1,24 @@ import ast import logging -from .core import TrackedContextTransformer, make_function_transformer, primitive_ast_types +from .core import TrackedContextTransformer, make_function_transformer, primitive_ast_types, iterable_ast_types log = logging.getLogger(__name__) # noinspection PyPep8Naming class CollapseTransformer(TrackedContextTransformer): + collapse_iterables = False + def visit_Name(self, node): res = self.resolve_literal(node) if isinstance(res, primitive_ast_types): return res + if isinstance(res, iterable_ast_types): + if self.collapse_iterables: + return res + else: + log.debug("Not collapsing iterable {}. Change this setting with collapse_literals(collapse_iterables=True)".format(res)) return node def visit_BinOp(self, node): diff --git a/pragma/core/resolve/__init__.py b/pragma/core/resolve/__init__.py index 567eb7c..2c88e09 100644 --- a/pragma/core/resolve/__init__.py +++ b/pragma/core/resolve/__init__.py @@ -89,11 +89,13 @@ def resolve_name_or_attribute(node, ctxt): float_types = (float,) primitive_types = tuple([str, bytes, bool, type(None)] + list(num_types) + list(float_types)) +iterable_types = (list, tuple) try: primitive_ast_types = (ast.Num, ast.Str, ast.Bytes, ast.NameConstant, ast.Constant, ast.JoinedStr) except AttributeError: # Python <3.6 primitive_ast_types = (ast.Num, ast.Str, ast.Bytes, ast.NameConstant) +iterable_ast_types = (ast.List, ast.Tuple) def make_binop(op): diff --git a/pragma/core/transformer.py b/pragma/core/transformer.py index 348b79d..1b30dbf 100644 --- a/pragma/core/transformer.py +++ b/pragma/core/transformer.py @@ -347,7 +347,7 @@ def visit_ExceptHandler(self, node): def make_function_transformer(transformer_type, name, description, **transformer_kwargs): @optional_argument_decorator @magic_contract - def transform(return_source=False, save_source=True, function_globals=None, **kwargs): + def transform(return_source=False, save_source=True, function_globals=None, collapse_iterables=False, **kwargs): """ :param return_source: Returns the transformed function's source code instead of compiling it :type return_source: bool @@ -355,6 +355,8 @@ def transform(return_source=False, save_source=True, function_globals=None, **kw :type save_source: bool :param function_globals: Overridden global name assignments to use when processing the function :type function_globals: dict|None + :param collapse_iterables: Collapse iterable types + :type collapse_iterables: bool :param kwargs: Any other environmental variables to provide during unrolling :type kwargs: dict :return: The transformed function, or its source code if requested @@ -374,6 +376,7 @@ def inner(f): glbls.update(function_globals) # print({k: v for k, v in glbls.items() if k not in globals()}) trans = transformer_type(DictStack(glbls, kwargs), **transformer_kwargs) + trans.collapse_iterables = collapse_iterables f_mod.body[0].decorator_list = [] f_mod = trans.visit(f_mod) # print(astor.dump_tree(f_mod)) diff --git a/tests/pytest.ini b/tests/pytest.ini new file mode 100644 index 0000000..397ef63 --- /dev/null +++ b/tests/pytest.ini @@ -0,0 +1,4 @@ +[pytest] +filterwarnings = + ignore::DeprecationWarning +addopts = -s --log-cli-level 30 \ No newline at end of file diff --git a/tests/test_collapse_literals.py b/tests/test_collapse_literals.py index 6238c4f..0efcf31 100644 --- a/tests/test_collapse_literals.py +++ b/tests/test_collapse_literals.py @@ -356,6 +356,20 @@ def f(): self.assertSourceEqual(f, result) + def test_iterable_option(self): + a = [1, 2, 3, 4] + + @pragma.collapse_literals(collapse_iterables=True) + def f(): + x = a + + result = ''' + def f(): + x = [1, 2, 3, 4] + ''' + + self.assertSourceEqual(f, result) + def test_reduction(self): a = [1, 2, 3] From 681c5af95d4015a42fa970304a7ab8f3298f2f5c Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 13 Oct 2020 04:15:30 -0600 Subject: [PATCH 2/3] Explicit collapse option for fine tuning in a namespace with lots of globals --- pragma/core/transformer.py | 20 ++++++++++++++------ tests/test_collapse_literals.py | 15 +++++++++++++++ 2 files changed, 29 insertions(+), 6 deletions(-) diff --git a/pragma/core/transformer.py b/pragma/core/transformer.py index 1b30dbf..5fe747e 100644 --- a/pragma/core/transformer.py +++ b/pragma/core/transformer.py @@ -347,7 +347,7 @@ def visit_ExceptHandler(self, node): def make_function_transformer(transformer_type, name, description, **transformer_kwargs): @optional_argument_decorator @magic_contract - def transform(return_source=False, save_source=True, function_globals=None, collapse_iterables=False, **kwargs): + def transform(return_source=False, save_source=True, function_globals=None, collapse_iterables=False, explicit_only=False, **kwargs): """ :param return_source: Returns the transformed function's source code instead of compiling it :type return_source: bool @@ -357,6 +357,8 @@ def transform(return_source=False, save_source=True, function_globals=None, coll :type function_globals: dict|None :param collapse_iterables: Collapse iterable types :type collapse_iterables: bool + :param explicit_only: Whether to use global variables or just keyword and function_globals in the replacement context + :type explicit_only: bool :param kwargs: Any other environmental variables to provide during unrolling :type kwargs: dict :return: The transformed function, or its source code if requested @@ -366,11 +368,17 @@ def transform(return_source=False, save_source=True, function_globals=None, coll @magic_contract(f='Callable', returns='Callable|str') def inner(f): f_mod, f_body, f_file = function_ast(f) - # Grab function globals - glbls = f.__globals__.copy() - # Grab function closure variables - if isinstance(f.__closure__, tuple): - glbls.update({k: v.cell_contents for k, v in zip(f.__code__.co_freevars, f.__closure__)}) + if not explicit_only: + # Grab function globals + glbls = f.__globals__.copy() + # Grab function closure variables + if isinstance(f.__closure__, tuple): + glbls.update({k: v.cell_contents for k, v in zip(f.__code__.co_freevars, f.__closure__)}) + else: + # Initialize empty context + if function_globals is None and len(kwargs) == 0: + log.warning("No global context nor function context. No collapse will occur") + glbls = dict() # Apply manual globals override if function_globals is not None: glbls.update(function_globals) diff --git a/tests/test_collapse_literals.py b/tests/test_collapse_literals.py index 0efcf31..223f244 100644 --- a/tests/test_collapse_literals.py +++ b/tests/test_collapse_literals.py @@ -437,3 +437,18 @@ def f(): self.assertSourceEqual(f, result) self.assertEqual(f(), 4) + + + def test_explicit_collapse(self): + a = 2 + b = 3 + @pragma.collapse_literals(explicit_only=True, b=b) + def f(): + x = a + y = b + result = ''' + def f(): + x = a + y = 3 + ''' + self.assertSourceEqual(f, result) From 1ae9ec860d2eb394ac0ede9ebc33cb780c6d7dd1 Mon Sep 17 00:00:00 2001 From: Alex Tait Date: Sat, 17 Oct 2020 03:00:03 -0600 Subject: [PATCH 3/3] trivial test of no collapsing --- tests/test_collapse_literals.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/test_collapse_literals.py b/tests/test_collapse_literals.py index eb32531..10eb008 100644 --- a/tests/test_collapse_literals.py +++ b/tests/test_collapse_literals.py @@ -497,3 +497,12 @@ def f(): y = 3 ''' self.assertSourceEqual(f, result) + + @pragma.collapse_literals(explicit_only=True) + def f(): + x = a + result = ''' + def f(): + x = a + ''' + self.assertSourceEqual(f, result)