Support declaring compatible PyTeal version in source code (#429)

jdtzmn · web-flow · commit cf560841b94a · 2022-07-06T16:01:05.000-07:00
* Add semantic version dependency

* Add `pragma` to enforce compatible compiler versions

* Raise `TealInternalError` instead of `TealInputError`

* Switch back to `pkg_resources` and convert PEP 440 to semantic version 2.0.0

* Fix linter errors

* Make pep 440 converter private

* Add `Pragma` expression

* Refactor pragma functions from compiler to ast to avoid circular dependencies
* Add `is_valid_compiler_version` check helper for the `Pragma` expression
* Use `TealProgramError` instead of `TealInternalError`

* Test underlying TEAL is unchanged

* Refactor underlying pragma methods to parent directory

* Inherit from `Expr` instead

* Fix unclear docstring

* Document init

Document init

* Add caret compiler version tests

* Ignore unused fixture imports
diff --git a/mypy.ini b/mypy.ini
@@ -1,5 +1,8 @@
 [mypy]
 
+[mypy-semantic_version.*]
+ignore_missing_imports = True
+
 [mypy-pytest.*]
 ignore_missing_imports = True
 
diff --git a/pyteal/__init__.py b/pyteal/__init__.py
@@ -1,5 +1,6 @@
 from pyteal.ast import *
 from pyteal.ast import __all__ as ast_all
+from pyteal.pragma import pragma
 from pyteal.ir import *
 from pyteal.ir import __all__ as ir_all
 from pyteal.compiler import (
@@ -16,6 +17,7 @@
     TealTypeError,
     TealInputError,
     TealCompileError,
+    TealPragmaError,
 )
 from pyteal.config import (
     MAX_GROUP_SIZE,
@@ -33,13 +35,15 @@
         "MIN_TEAL_VERSION",
         "DEFAULT_TEAL_VERSION",
         "CompileOptions",
+        "pragma",
         "compileTeal",
         "OptimizeOptions",
         "TealType",
         "TealInternalError",
         "TealTypeError",
         "TealInputError",
         "TealCompileError",
+        "TealPragmaError",
         "MAX_GROUP_SIZE",
         "NUM_SLOTS",
         "RETURN_HASH_PREFIX",
diff --git a/pyteal/__init__.pyi b/pyteal/__init__.pyi
@@ -3,6 +3,7 @@
 
 from pyteal.ast import *
 from pyteal.ast import __all__ as ast_all
+from pyteal.pragma import pragma
 from pyteal.ir import *
 from pyteal.ir import __all__ as ir_all
 from pyteal.compiler import (
@@ -19,6 +20,7 @@ from pyteal.errors import (
     TealTypeError,
     TealInputError,
     TealCompileError,
+    TealPragmaError,
 )
 from pyteal.config import (
     MAX_GROUP_SIZE,
@@ -152,6 +154,7 @@ __all__ = [
     "OptimizeOptions",
     "Or",
     "Pop",
+    "Pragma",
     "RETURN_HASH_PREFIX",
     "Reject",
     "Return",
@@ -185,6 +188,7 @@ __all__ = [
     "TealInternalError",
     "TealLabel",
     "TealOp",
+    "TealPragmaError",
     "TealSimpleBlock",
     "TealType",
     "TealTypeError",
@@ -202,4 +206,5 @@ __all__ = [
     "WideRatio",
     "abi",
     "compileTeal",
+    "pragma",
 ]
diff --git a/pyteal/ast/__init__.py b/pyteal/ast/__init__.py
@@ -40,6 +40,7 @@
 from pyteal.ast.array import Array
 from pyteal.ast.tmpl import Tmpl
 from pyteal.ast.nonce import Nonce
+from pyteal.ast.pragma import Pragma
 
 # unary ops
 from pyteal.ast.unaryexpr import (
@@ -199,6 +200,7 @@
     "Array",
     "Tmpl",
     "Nonce",
+    "Pragma",
     "UnaryExpr",
     "Btoi",
     "Itob",
diff --git a/pyteal/ast/pragma.py b/pyteal/ast/pragma.py
@@ -0,0 +1,50 @@
+from typing import TYPE_CHECKING, Any
+
+from pyteal.ast.expr import Expr
+from pyteal.pragma import is_valid_compiler_version, pragma
+
+if TYPE_CHECKING:
+    from pyteal.compiler import CompileOptions
+
+
+class Pragma(Expr):
+    """A meta expression which defines a pragma for a specific subsection of PyTeal code.
+
+    This expression does not affect the underlying compiled TEAL code in any way."""
+
+    def __init__(self, child: Expr, *, compiler_version: str, **kwargs: Any) -> None:
+        """Define a pragma for a specific subsection of PyTeal code.
+
+        The Pragma expression does not affect the underlying compiled TEAL code in any way,
+        it merely sets a pragma for the underlying expression.
+
+        Args:
+            child: The expression to wrap.
+            compiler_version: Acceptable versions of the compiler. Will fail if the current PyTeal version
+                is not contained in the range. Follows the npm `semver range scheme <https://github.com/npm/node-semver#ranges>`_
+                for specifying compatible versions.
+        """
+        super().__init__()
+
+        self.child = child
+
+        if not is_valid_compiler_version(compiler_version):
+            raise ValueError("Invalid compiler version: {}".format(compiler_version))
+        self.compiler_version = compiler_version
+
+    def __teal__(self, options: "CompileOptions"):
+        pragma(compiler_version=self.compiler_version)
+
+        return self.child.__teal__(options)
+
+    def __str__(self):
+        return "(pragma {})".format(self.child)
+
+    def type_of(self):
+        return self.child.type_of()
+
+    def has_return(self):
+        return self.child.has_return()
+
+
+Pragma.__module__ = "pyteal"
diff --git a/pyteal/ast/pragma_test.py b/pyteal/ast/pragma_test.py
@@ -0,0 +1,59 @@
+import pytest
+from tests.mock_version import mock_version  # noqa: F401
+
+import pyteal as pt
+
+
+@pytest.mark.usefixtures("mock_version")
+@pytest.mark.parametrize(
+    "version, compiler_version, should_error",
+    [
+        # valid
+        ("0.12.0", "0.12.0", False),
+        (
+            "1.0.0+AVM7.1",
+            "=1.0.0",
+            False,
+        ),
+        # invalid
+        ("0.13.0", "0.13.1", True),
+        ("1.2.3a2", "<0.8.0 || >=0.12.0", True),
+    ],
+)
+def test_pragma_expr(compiler_version, should_error):
+    program = pt.Pragma(pt.Approve(), compiler_version=compiler_version)
+
+    if should_error:
+        with pytest.raises(pt.TealPragmaError):
+            pt.compileTeal(program, mode=pt.Mode.Application, version=6)
+    else:
+        pt.compileTeal(program, mode=pt.Mode.Application, version=6)
+
+
+def test_pragma_expr_does_not_change():
+    without_pragma = pt.Seq(pt.Pop(pt.Add(pt.Int(1), pt.Int(2))), pt.Return(pt.Int(1)))
+    pragma = pt.Pragma(without_pragma, compiler_version=">=0.0.0")
+
+    compiled_with_pragma = pt.compileTeal(pragma, mode=pt.Mode.Application, version=6)
+    compiled_without_pragma = pt.compileTeal(
+        without_pragma, mode=pt.Mode.Application, version=6
+    )
+
+    assert compiled_with_pragma == compiled_without_pragma
+
+
+def test_pragma_expr_has_return():
+    exprWithReturn = pt.Pragma(pt.Return(pt.Int(1)), compiler_version=">=0.0.0")
+    assert exprWithReturn.has_return()
+
+    exprWithoutReturn = pt.Pragma(pt.Int(1), compiler_version=">=0.0.0")
+    assert not exprWithoutReturn.has_return()
+
+
+@pytest.mark.parametrize(
+    "compiler_version",
+    ["not a version", ">=0.1.1,<0.3.0", "1.2.3aq"],  # incorrect spec  # invalid PEP 440
+)
+def test_pragma_expr_invalid_compiler_version(compiler_version):
+    with pytest.raises(ValueError):
+        pt.Pragma(pt.Approve(), compiler_version=compiler_version)
diff --git a/pyteal/errors.py b/pyteal/errors.py
@@ -64,6 +64,17 @@ def __eq__(self, other) -> bool:
 TealCompileError.__module__ = "pyteal"
 
 
+class TealPragmaError(Exception):
+    def __init__(self, message: str) -> None:
+        self.message = message
+
+    def __str__(self):
+        return self.message
+
+
+TealPragmaError.__module__ = "pyteal"
+
+
 def verifyTealVersion(minVersion: int, version: int, msg: str):
     if minVersion > version:
         msg = "{}. Minimum version needed is {}, but current version being compiled is {}".format(
diff --git a/pyteal/pragma/__init__.py b/pyteal/pragma/__init__.py
@@ -0,0 +1,6 @@
+from pyteal.pragma.pragma import is_valid_compiler_version, pragma
+
+__all__ = [
+    "is_valid_compiler_version",
+    "pragma",
+]
diff --git a/pyteal/pragma/pragma.py b/pyteal/pragma/pragma.py
@@ -0,0 +1,120 @@
+import re
+import pkg_resources
+from typing import Any
+from semantic_version import Version, NpmSpec
+
+from pyteal.errors import TealPragmaError
+
+
+def __convert_pep440_compiler_version(compiler_version: str):
+    """Convert PEP 440 version identifiers to valid NPM versions.
+
+    For example:
+        "1.0.0" -> "1.0.0"
+        "1.0.0a1" -> "1.0.0-a1"
+        "<0.5.0+local || >=1.0.0a9.post1.dev2" -> "<0.5.0 || >=1.0.0-alpha9.1.2"
+    """
+    NUMBER = r"(?:x|X|\*|0|[1-9][0-9]*)"
+    LOCAL = r"[a-zA-Z0-9.]*"
+    TRIM_PREFIX_RE = re.compile(
+        r"""
+            (?:v)?                                          # Strip optional initial v
+            (?P<op><|<=|>=|>|=|\^|~|)                       # Operator, can be empty
+            (?P<major>{nb})(?:\.(?P<minor>{nb})(?:\.(?P<patch>{nb}))?)?
+            (?:(?P<prerel_type>a|b|rc)(?P<prerel>{nb}))?    # Optional pre-release
+            (?:\.post(?P<postrel>{nb}))?                    # Optional post-release
+            (?:\.dev(?P<dev>{nb}))?                         # Optional dev release
+            (?:\+(?P<local>{lcl}))?                         # Optional local version
+        """.format(
+            nb=NUMBER,
+            lcl=LOCAL,
+        ),
+        re.VERBOSE,
+    )
+
+    def match_replacer(match: re.Match):
+        (
+            op,
+            major,
+            minor,
+            patch,
+            prerel_type,
+            prerel,
+            postrel,
+            dev,
+            local,
+        ) = match.groups()
+
+        # Base version (major/minor/patch)
+        base_version = "{}.{}.{}".format(major or "0", minor or "0", patch or "0")
+
+        # Combine prerel, postrel, and dev
+        combined_additions = []
+        short_prerel_type_to_long = {
+            "a": "alpha",
+            "b": "beta",
+            "rc": "rc",
+        }
+        if prerel_type is not None:
+            combined_additions.append(short_prerel_type_to_long[prerel_type] + prerel)
+        if len(combined_additions) > 0 or postrel is not None or dev is not None:
+            combined_additions.append(postrel or "0")
+        if len(combined_additions) > 0 or dev is not None:
+            combined_additions.append(dev or "0")
+        combined_additions_str = ".".join(combined_additions)
+
+        # Build full_version
+        full_version = base_version
+        if len(combined_additions) > 0:
+            full_version += "-" + combined_additions_str
+        if local is not None:
+            full_version += "+" + local.lower()
+
+        if op is not None:
+            return op + full_version
+        return full_version
+
+    return re.sub(TRIM_PREFIX_RE, match_replacer, compiler_version)
+
+
+def is_valid_compiler_version(compiler_version: str):
+    """Check if the compiler version is valid.
+
+    Args:
+        compiler_version: The compiler version to check.
+
+    Returns:
+        True if the compiler version is a valid NPM specification range
+        using either the PEP 440 or semantic version format, otherwise False.
+    """
+    try:
+        pep440_converted = __convert_pep440_compiler_version(compiler_version)
+        NpmSpec(pep440_converted)
+        return True
+    except ValueError:
+        return False
+
+
+def pragma(
+    *,
+    compiler_version: str,
+    **kwargs: Any,
+) -> None:
+    """
+    Specify pragmas for the compiler.
+
+    Args:
+        compiler_version: Acceptable versions of the compiler. Will fail if the current PyTeal version
+            is not contained in the range. Follows the npm `semver range scheme <https://github.com/npm/node-semver#ranges>`_
+            for specifying compatible versions.
+    """
+    pkg_version = pkg_resources.require("pyteal")[0].version
+    pyteal_version = Version(__convert_pep440_compiler_version(pkg_version))
+    if pyteal_version not in NpmSpec(
+        __convert_pep440_compiler_version(compiler_version)
+    ):
+        raise TealPragmaError(
+            "PyTeal version {} is not compatible with compiler version {}".format(
+                pkg_version, compiler_version
+            )
+        )
diff --git a/pyteal/pragma/pragma_test.py b/pyteal/pragma/pragma_test.py
diff --git a/requirements.txt b/requirements.txt
diff --git a/setup.py b/setup.py
diff --git a/tests/mock_version.py b/tests/mock_version.py
