Make pylance recognize wildcard imports

.travis.yml

@@ -6,6 +6,7 @@ install:
 script:
   - pytest
   - mypy pyteal
+  - python3 -c "import pyteal" scripts/generate_init.py --check
   - black --check .
 
 jobs:

pyteal/__init__.py

@@ -13,6 +13,7 @@
 from .errors import TealInternalError, TealTypeError, TealInputError, TealCompileError
 from .config import MAX_GROUP_SIZE, NUM_SLOTS
 
+# begin __all__
 __all__ = (
     ast_all
     + ir_all
@@ -31,3 +32,4 @@
         "NUM_SLOTS",
     ]
 )
+# end __all__

pyteal/__init__.pyi

@@ -0,0 +1,161 @@
+## File generated from scripts/generate_init.py.
+## DO NOT EDIT DIRECTLY
+
+from .ast import *
+from .ast import __all__ as ast_all
+from .ir import *
+from .ir import __all__ as ir_all
+from .compiler import (
+    MAX_TEAL_VERSION,
+    MIN_TEAL_VERSION,
+    DEFAULT_TEAL_VERSION,
+    CompileOptions,
+    compileTeal,
+)
+from .types import TealType
+from .errors import TealInternalError, TealTypeError, TealInputError, TealCompileError
+from .config import MAX_GROUP_SIZE, NUM_SLOTS
+
+__all__ = [
+    "Expr",
+    "LeafExpr",
+    "Addr",
+    "Bytes",
+    "Int",
+    "EnumInt",
+    "Arg",
+    "TxnType",
+    "TxnField",
+    "TxnExpr",
+    "TxnaExpr",
+    "TxnArray",
+    "TxnObject",
+    "Txn",
+    "GtxnExpr",
+    "GtxnaExpr",
+    "TxnGroup",
+    "Gtxn",
+    "GeneratedID",
+    "ImportScratchValue",
+    "Global",
+    "GlobalField",
+    "App",
+    "AppField",
+    "OnComplete",
+    "AppParam",
+    "AssetHolding",
+    "AssetParam",
+    "InnerTxnBuilder",
+    "InnerTxn",
+    "Array",
+    "Tmpl",
+    "Nonce",
+    "UnaryExpr",
+    "Btoi",
+    "Itob",
+    "Len",
+    "BitLen",
+    "Sha256",
+    "Sha512_256",
+    "Keccak256",
+    "Not",
+    "BitwiseNot",
+    "Sqrt",
+    "Pop",
+    "Balance",
+    "MinBalance",
+    "BinaryExpr",
+    "Add",
+    "Minus",
+    "Mul",
+    "Div",
+    "Mod",
+    "Exp",
+    "BitwiseAnd",
+    "BitwiseOr",
+    "BitwiseXor",
+    "ShiftLeft",
+    "ShiftRight",
+    "Eq",
+    "Neq",
+    "Lt",
+    "Le",
+    "Gt",
+    "Ge",
+    "GetBit",
+    "GetByte",
+    "Ed25519Verify",
+    "Substring",
+    "Extract",
+    "Suffix",
+    "SetBit",
+    "SetByte",
+    "NaryExpr",
+    "And",
+    "Or",
+    "Concat",
+    "WideRatio",
+    "If",
+    "Cond",
+    "Seq",
+    "Assert",
+    "Err",
+    "Return",
+    "Approve",
+    "Reject",
+    "Subroutine",
+    "SubroutineDefinition",
+    "SubroutineDeclaration",
+    "SubroutineCall",
+    "ScratchSlot",
+    "ScratchLoad",
+    "ScratchStore",
+    "ScratchStackStore",
+    "ScratchVar",
+    "MaybeValue",
+    "BytesAdd",
+    "BytesMinus",
+    "BytesDiv",
+    "BytesMul",
+    "BytesMod",
+    "BytesAnd",
+    "BytesOr",
+    "BytesXor",
+    "BytesEq",
+    "BytesNeq",
+    "BytesLt",
+    "BytesLe",
+    "BytesGt",
+    "BytesGe",
+    "BytesNot",
+    "BytesZero",
+    "ExtractUint16",
+    "ExtractUint32",
+    "ExtractUint64",
+    "Log",
+    "While",
+    "For",
+    "Break",
+    "Continue",
+    "Op",
+    "Mode",
+    "TealComponent",
+    "TealOp",
+    "TealLabel",
+    "TealBlock",
+    "TealSimpleBlock",
+    "TealConditionalBlock",
+    "LabelReference",
+    "MAX_TEAL_VERSION",
+    "MIN_TEAL_VERSION",
+    "DEFAULT_TEAL_VERSION",
+    "CompileOptions",
+    "compileTeal",
+    "TealType",
+    "TealInternalError",
+    "TealTypeError",
+    "TealInputError",
+    "TealCompileError",
+    "MAX_GROUP_SIZE",
+    "NUM_SLOTS",
+]

scripts/generate_init.py

@@ -0,0 +1,88 @@
+import argparse, os, sys
+from pyteal import __all__ as static_all
+
+
+# Start of the template to be appended to
+pyi_template = """## File generated from scripts/generate_init.py.
+## DO NOT EDIT DIRECTLY
+
+"""
+
+# Template for __all__ export list
+all_template = """__all__ = [
+    {},
+]"""
+
+# Flags to denote the beginning/end of the __all__ exports in __init__.py
+begin_flag = "# begin __all__"
+end_flag = "# end __all__"
+
+# Make it safe to run from anywhere
+curr_dir = os.path.dirname(os.path.abspath(__file__))
+orig_dir = os.path.join(curr_dir, os.path.join("..", "pyteal"))
+
+# Path to pyi
+pyi_file = "__init__.pyi"
+orig_file = os.path.join(orig_dir, pyi_file)
+
+# Path to py
+py_file = "__init__.py"
+init_file = os.path.join(orig_dir, py_file)
+
+
+def generate_tmp():
+    with open(init_file, "r") as f:
+        init_contents = f.read()
+
+    start_idx = init_contents.index(begin_flag)
+    end_idx = init_contents.index(end_flag)
+
+    all_imports = ",\n    ".join(['"{}"'.format(s) for s in static_all])
+
+    return (
+        pyi_template
+        + init_contents[:start_idx]
+        + all_template.format(all_imports)
+        + init_contents[end_idx + len(end_flag) :]
+    )
+
+
+def is_different(regen):
+    if not os.path.exists(orig_file):
+        return True
+
+    with open(orig_file, "r") as f:
+        orig_lines = f.readlines()
+
+    curr_lines = regen.split("\n")
+
+    return orig_lines == curr_lines
+
+
+def overwrite(regen):
+    with open(orig_file, "w") as f:
+        f.write(regen)
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--check",
+        action="store_true",
+        help="Only check if the generated file would change",
+    )
+    args = parser.parse_args()
+
+    regen = generate_tmp()
+
+    if args.check:
+        if is_different(regen):
+            print(
+                "The __init__.pyi needs to be regenerated. Please run scripts/generate_init.py"
+            )
+            sys.exit(1)
+        print("No changes in __init__.py")
+        sys.exit(0)
+
+    overwrite(regen)