Skip to content

Commit

Permalink
Make pylance recognize wildcard imports (#133)
Browse files Browse the repository at this point in the history
* adding exports directly to top level __all__

* apply black formatter

* adding initial generate script

* fmt

* rm all from all

* adding check to travis

* reading in original __init__ and using its imports, dont write to filesystem if --check is passed

* make messages more profesh

* fix flags after black formatted them

* y

* flippin black formatter

* help text fix

* asdfasdf
  • Loading branch information
barnjamin authored Nov 5, 2021
1 parent e30846a commit fdf003d
Show file tree
Hide file tree
Showing 4 changed files with 252 additions and 0 deletions.
1 change: 1 addition & 0 deletions .travis.yml
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ install:
script:
- pytest
- mypy pyteal
- python3 -c "import pyteal" scripts/generate_init.py --check
- black --check .

jobs:
Expand Down
2 changes: 2 additions & 0 deletions pyteal/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -31,3 +32,4 @@
"NUM_SLOTS",
]
)
# end __all__
161 changes: 161 additions & 0 deletions pyteal/__init__.pyi
Original file line number Diff line number Diff line change
@@ -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",
]
88 changes: 88 additions & 0 deletions scripts/generate_init.py
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit fdf003d

Please sign in to comment.