chore: pr feedback

ApeWorX · Apr 11, 2024 · 6bc6b26 · 6bc6b26
1 parent 8f98411
commit 6bc6b26
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 26 deletions.
diff --git a/ape_vyper/__init__.py b/ape_vyper/__init__.py
@@ -1,7 +1,6 @@
 from ape import plugins
 
-from ._utils import Extension
-from .compiler import VyperCompiler, VyperConfig
+from .compiler import FileType, VyperCompiler, VyperConfig
 
 
 @plugins.register(plugins.Config)
@@ -11,4 +10,4 @@ def config_class():
 
 @plugins.register(plugins.CompilerPlugin)
 def register_compiler():
-    return tuple([e.value for e in Extension]), VyperCompiler
+    return tuple(e.value for e in FileType), VyperCompiler
diff --git a/ape_vyper/_utils.py b/ape_vyper/_utils.py
diff --git a/ape_vyper/compiler.py b/ape_vyper/compiler.py
@@ -3,6 +3,7 @@
 import shutil
 import time
 from base64 import b64encode
+from enum import Enum
 from fnmatch import fnmatch
 from importlib import import_module
 from pathlib import Path
@@ -29,7 +30,6 @@
 from vvm import compile_standard as vvm_compile_standard
 from vvm.exceptions import VyperError  # type: ignore
 
-from ape_vyper._utils import Extension
 from ape_vyper.ast import source_to_abi
 from ape_vyper.exceptions import (
     RUNTIME_ERROR_MAP,
@@ -55,6 +55,14 @@
 Optimization = Union[str, bool]
 
 
+class FileType(str, Enum):
+    SOURCE = ".vy"
+    INTERFACE = ".vyi"
+
+    def __str__(self) -> str:
+        return self.value
+
+
 class VyperConfig(PluginConfig):
     version: Optional[SpecifierSet] = None
     """
@@ -241,12 +249,12 @@ def get_imports(
 
                 # NOTE: Defaults to JSON (assuming from input JSON or a local JSON),
                 #  unless a Vyper file exists.
-                if (base_path / f"{prefix}{Extension.VY.value}").is_file():
-                    ext = Extension.VY.value
-                elif (base_path / f"{prefix}{Extension.VY.value}").is_file():
-                    ext = Extension.VY.value
-                elif (base_path / f"{prefix}{Extension.VYI.value}").is_file():
-                    ext = Extension.VYI.value
+                if (base_path / f"{prefix}{FileType.SOURCE}").is_file():
+                    ext = FileType.SOURCE.value
+                elif (base_path / f"{prefix}{FileType.SOURCE}").is_file():
+                    ext = FileType.INTERFACE.value
+                elif (base_path / f"{prefix}{FileType.INTERFACE}").is_file():
+                    ext = FileType.INTERFACE.value
                 else:
                     ext = ".json"
 
@@ -765,7 +773,7 @@ def get_compiler_settings(
         valid_paths = [
             p
             for p in contract_filepaths
-            if get_full_extension(p) == Extension.VY
+            if get_full_extension(p) == FileType.SOURCE
             and not str(p).startswith(str(contracts_path / "interfaces"))
         ]
         files_by_vyper_version = self.get_version_map(valid_paths, base_path=contracts_path)
@@ -1082,7 +1090,7 @@ def _get_traceback(
                 called_contract, sub_calldata = self._create_contract_from_call(frame)
                 if called_contract:
                     ext = get_full_extension(Path(called_contract.source_id))
-                    if ext in [e.value for e in Extension]:
+                    if ext in [e.value for e in FileType]:
                         # Called another Vyper contract.
                         sub_trace = self._get_traceback(
                             called_contract, trace, sub_calldata, previous_depth=frame.depth
@@ -1528,11 +1536,3 @@ def _is_fallback_check(opcodes: List[str], op: str) -> bool:
         and opcodes[6] == "SHR"
         and opcodes[5] == "0xE0"
     )
-
-
-# def _version_to_specifier(version: str) -> str:
-#     pragma_str = " ".join(version.split()).replace("^", "~=")
-#     if pragma_str and pragma_str[0].isnumeric():
-#         return f"=={pragma_str}"
-#
-#     return pragma_str