Merge pull request #88 from dflook/annotations

Make remove annotations transform configurable

Author: dflook

.github/workflows/release_test.yaml

@@ -112,9 +112,9 @@ jobs:
       - name: Test typing
         run: |
           pip3.11 install mypy types-setuptools
-          mypy typing_test/test_typing.py
+          mypy --strict typing_test/test_typing.py
 
-          if mypy typing_test/test_badtyping.py; then
+          if mypy --strict typing_test/test_badtyping.py; then
             echo "Bad types weren't detected"
             exit 1
           fi

docs/source/api_usage.rst

@@ -2,4 +2,9 @@ Package Reference
 =================
 
 .. automodule:: python_minifier
-   :members: minify, awslambda, unparse, UnstableMinification
+
+.. autofunction:: minify
+.. autoclass:: RemoveAnnotationsOptions
+.. autofunction:: awslambda
+.. autofunction:: unparse
+.. autoclass:: UnstableMinification

docs/source/conf.py

@@ -89,7 +89,7 @@ def create_example(option):
 #
 # This is also used if you do content translation via gettext catalogs.
 # Usually you set "language" from the command line for these cases.
-language = None
+language = 'en'
 
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
@@ -116,7 +116,7 @@ def create_example(option):
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ['_static']
+html_static_path = []
 
 # Custom sidebar templates, must be a dictionary that maps document names
 # to template names.

docs/source/transforms/remove_annotations.py

@@ -1,14 +1,6 @@
-def compile(source: "something compilable",
-            filename: "where the compilable thing comes from",
-            mode: "A string annotation"):
-    a: int = 1
-    b: str
-
-    def inner():
-        nonlocal b
-        b = 'hello'
-
-    inner()
-    print(b)
-
-compile(None, None, None)
+class A:
+    b: int
+    c: int=2
+    def a(self, val: str) -> None:
+        b: int
+        c: int=2

docs/source/transforms/remove_annotations.rst

@@ -1,24 +1,47 @@
 Remove Annotations
 ==================
 
-This transform removes function annotations and variable annotations.
-
-This transform is generally safe to use. Although the annotations have no meaning to the python language,
+This transform removes annotations. Although the annotations have no meaning to the python language,
 they are made available at runtime. Some python library features require annotations to be kept.
 
-If these are detected, annotations are kept for that class:
+Annotations can be removed from:
+
+    - Function arguments
+    - Function return
+    - Variables
+    - Class attributes
+
+By default annotations are removed from variables, function arguments and function return, but not from class attributes.
+
+This transform is generally safe to use with the default options. If you know the module requires the annotations to be kept, disable this transform.
+Class attribute annotations can often be used by other modules, so it is recommended to keep them unless you know they are not used.
+
+When removing class attribute annotations is enabled, annotations are kept for classes that are derived from:
 
     - dataclasses.dataclass
     - typing.NamedTuple
     - typing.TypedDict
 
-If you know the module requires the annotations to be kept, disable this transform.
-
 If a variable annotation without assignment is used the annotation is changed to a literal zero instead of being removed.
 
-The transform is enabled by default.
-Disable this source transformation by passing the ``remove_annotations=False`` argument to the :func:`python_minifier.minify` function,
-or passing ``--no-remove-annotations`` to the pyminify command.
+Options
+-------
+
+These arguments can be used with the pyminify command:
+
+``--no-remove-variable-annotations`` disables removing variable annotations
+
+``--no-remove-return-annotations`` disables removing function return annotations
+
+``--no-remove-argument-annotations`` disables removing function argument annotations
+
+``--remove-class-attribute-annotations`` enables removing class attribute annotations
+
+``--no-remove-annotations`` disables removing all annotations, this transform will not do anything.
+
+When using the :func:`python_minifier.minify` function you can use the  ``remove_annotations`` argument to control this transform.
+You can pass a boolean ``True`` to remove all annotations or a boolean ``False`` to keep all annotations.
+You can also pass a :class:`python_minifier.RemoveAnnotationsOptions` instance to specify which annotations to remove.
 
 Example
 -------

setup.py

@@ -23,7 +23,7 @@
     use_scm_version=True,
     package_dir={'': 'src'},
     packages=find_packages('src'),
-    package_data={"python_minifier": ["py.typed", "__init__.pyi"]},
+    package_data={"python_minifier": ["py.typed", "*.pyi", "rename/*.pyi", "transforms/*.pyi"]},
     long_description=long_desc,
     long_description_content_type='text/markdown',
 

src/python_minifier/__init__.py

@@ -21,6 +21,7 @@
 
 from python_minifier.transforms.combine_imports import CombineImports
 from python_minifier.transforms.remove_annotations import RemoveAnnotations
+from python_minifier.transforms.remove_annotations_options import RemoveAnnotationsOptions
 from python_minifier.transforms.remove_asserts import RemoveAsserts
 from python_minifier.transforms.remove_debug import RemoveDebug
 from python_minifier.transforms.remove_explicit_return_none import RemoveExplicitReturnNone
@@ -52,7 +53,7 @@ def __str__(self):
 def minify(
     source,
     filename=None,
-    remove_annotations=True,
+    remove_annotations=RemoveAnnotationsOptions(),
     remove_pass=True,
     remove_literal_statements=False,
     combine_imports=True,
@@ -80,7 +81,8 @@ def minify(
     :param str source: The python module source code
     :param str filename: The original source filename if known
 
-    :param bool remove_annotations: If type annotations should be removed where possible
+    :param remove_annotations: Configures the removal of type annotations. True removes all annotations, False removes none. RemoveAnnotationsOptions can be used to configure the removal of specific annotations.
+    :type remove_annotations: bool or RemoveAnnotationsOptions
     :param bool remove_pass: If Pass statements should be removed where possible
     :param bool remove_literal_statements: If statements consisting of a single literal should be removed, including docstrings
     :param bool combine_imports: Combine adjacent import statements where possible
@@ -115,8 +117,20 @@ def minify(
     if combine_imports:
         module = CombineImports()(module)
 
-    if remove_annotations:
-        module = RemoveAnnotations()(module)
+    if isinstance(remove_annotations, bool):
+        remove_annotations_options = RemoveAnnotationsOptions(
+            remove_variable_annotations=remove_annotations,
+            remove_return_annotations=remove_annotations,
+            remove_argument_annotations=remove_annotations,
+            remove_class_attribute_annotations=remove_annotations,
+        )
+    elif isinstance(remove_annotations, RemoveAnnotationsOptions):
+        remove_annotations_options = remove_annotations
+    else:
+        raise TypeError('remove_annotations must be a bool or RemoveAnnotationsOptions')
+
+    if remove_annotations_options:
+        module = RemoveAnnotations(remove_annotations_options)(module)
 
     if remove_pass:
         module = RemovePass()(module)

src/python_minifier/__init__.pyi

@@ -1,14 +1,15 @@
 import ast
-from typing import List, Text, AnyStr, Optional, Any
+from typing import List, Text, AnyStr, Optional, Any, Union
 
+from .transforms.remove_annotations_options import RemoveAnnotationsOptions as RemoveAnnotationsOptions
 
 class UnstableMinification(RuntimeError):
     def __init__(self, exception: Any, source: Any, minified: Any): ...
 
 def minify(
     source: AnyStr,
     filename: Optional[str] = ...,
-    remove_annotations: bool = ...,
+    remove_annotations: Union[bool, RemoveAnnotationsOptions] = ...,
     remove_pass: bool = ...,
     remove_literal_statements: bool = ...,
     combine_imports: bool = ...,

src/python_minifier/__main__.py

@@ -7,6 +7,7 @@
 from pkg_resources import get_distribution, DistributionNotFound
 
 from python_minifier import minify
+from python_minifier.transforms.remove_annotations_options import RemoveAnnotationsOptions
 
 try:
     version = get_distribution('python_minifier').version
@@ -113,12 +114,6 @@ def parse_args():
         help='Enable removing statements that are just a literal (including docstrings)',
         dest='remove_literal_statements',
     )
-    minification_options.add_argument(
-        '--no-remove-annotations',
-        action='store_false',
-        help='Disable removing function and variable annotations',
-        dest='remove_annotations',
-    )
     minification_options.add_argument(
         '--no-hoist-literals',
         action='store_false',
@@ -189,6 +184,39 @@ def parse_args():
         help='Replace explicit return None with a bare return',
         dest='remove_explicit_return_none',
     )
+
+    annotation_options = parser.add_argument_group('remove annotations options', 'Options that affect how annotations are removed')
+    annotation_options.add_argument(
+        '--no-remove-annotations',
+        action='store_false',
+        help='Disable removing all annotations',
+        dest='remove_annotations',
+    )
+    annotation_options.add_argument(
+        '--no-remove-variable-annotations',
+        action='store_false',
+        help='Disable removing variable annotations',
+        dest='remove_variable_annotations',
+    )
+    annotation_options.add_argument(
+        '--no-remove-return-annotations',
+        action='store_false',
+        help='Disable removing function return annotations',
+        dest='remove_return_annotations',
+    )
+    annotation_options.add_argument(
+        '--no-remove-argument-annotations',
+        action='store_false',
+        help='Disable removing function argument annotations',
+        dest='remove_argument_annotations',
+    )
+    annotation_options.add_argument(
+        '--remove-class-attribute-annotations',
+        action='store_true',
+        help='Enable removing class attribute annotations',
+        dest='remove_class_attribute_annotations',
+    )
+
     parser.add_argument('--version', '-v', action='version', version=version)
 
     args = parser.parse_args()
@@ -207,6 +235,10 @@ def parse_args():
         sys.stderr.write('error: path ' + args.path[0] + ' is a directory, --in-place required\n')
         sys.exit(1)
 
+    if args.remove_class_attribute_annotations and not args.remove_annotations:
+        sys.stderr.write('error: --remove-class-attribute-annotations would do nothing when used with --no-remove-annotations\n')
+        sys.exit(1)
+
     return args
 
 def source_modules(args):
@@ -237,12 +269,27 @@ def do_minify(source, filename, minification_args):
             names = [name.strip() for name in arg.split(',') if name]
             preserve_locals.extend(names)
 
+    if minification_args.remove_annotations is False:
+        remove_annotations = RemoveAnnotationsOptions(
+            remove_variable_annotations=False,
+            remove_return_annotations=False,
+            remove_argument_annotations=False,
+            remove_class_attribute_annotations=False,
+        )
+    else:
+        remove_annotations = RemoveAnnotationsOptions(
+            remove_variable_annotations=minification_args.remove_variable_annotations,
+            remove_return_annotations=minification_args.remove_return_annotations,
+            remove_argument_annotations=minification_args.remove_argument_annotations,
+            remove_class_attribute_annotations=minification_args.remove_class_attribute_annotations,
+        )
+
     return minify(
         source,
         filename=filename,
         combine_imports=minification_args.combine_imports,
         remove_pass=minification_args.remove_pass,
-        remove_annotations=minification_args.remove_annotations,
+        remove_annotations=remove_annotations,
         remove_literal_statements=minification_args.remove_literal_statements,
         hoist_literals=minification_args.hoist_literals,
         rename_locals=minification_args.rename_locals,

src/python_minifier/transforms/remove_annotations.py

@@ -1,6 +1,7 @@
 import ast
 import sys
 
+from python_minifier.transforms.remove_annotations_options import RemoveAnnotationsOptions
 from python_minifier.transforms.suite_transformer import SuiteTransformer
 
 
@@ -9,6 +10,11 @@ class RemoveAnnotations(SuiteTransformer):
     Remove type annotations from source
     """
 
+    def __init__(self, options):
+        assert isinstance(options, RemoveAnnotationsOptions)
+        self._options = options
+        super(RemoveAnnotations, self).__init__()
+
     def __call__(self, node):
         if sys.version_info < (3, 0):
             return node
@@ -19,7 +25,7 @@ def visit_FunctionDef(self, node):
         node.body = self.suite(node.body, parent=node)
         node.decorator_list = [self.visit(d) for d in node.decorator_list]
 
-        if hasattr(node, 'returns'):
+        if hasattr(node, 'returns') and self._options.remove_return_annotations:
             node.returns = None
 
         return node
@@ -37,21 +43,24 @@ def visit_arguments(self, node):
             node.kwonlyargs = [self.visit_arg(a) for a in node.kwonlyargs]
 
         if hasattr(node, 'varargannotation'):
-            node.varargannotation = None
+            if self._options.remove_argument_annotations:
+                node.varargannotation = None
         else:
             if node.vararg:
                 node.vararg = self.visit_arg(node.vararg)
 
         if hasattr(node, 'kwargannotation'):
-            node.kwargannotation = None
+            if self._options.remove_argument_annotations:
+                node.kwargannotation = None
         else:
             if node.kwarg:
                 node.kwarg = self.visit_arg(node.kwarg)
 
         return node
 
     def visit_arg(self, node):
-        node.annotation = None
+        if self._options.remove_argument_annotations:
+            node.annotation = None
         return node
 
     def visit_AnnAssign(self, node):
@@ -97,6 +106,14 @@ def is_typing_sensitive(node):
 
             return False
 
+        # is this a class attribute or a variable?
+        if isinstance(node.parent, ast.ClassDef):
+            if not self._options.remove_class_attribute_annotations:
+                return node
+        else:
+            if not self._options.remove_variable_annotations:
+                return node
+
         if is_dataclass_field(node) or is_typing_sensitive(node):
             return node
         elif node.value: