simplify a test

cyberthirst · Mar 19, 2024 · 6dfd08d · 6dfd08d
1 parent 039e5f4
commit 6dfd08d
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 7 deletions.
diff --git a/tests/functional/codegen/modules/test_flag_imports.py b/tests/functional/codegen/modules/test_flag_imports.py
@@ -1,6 +1,6 @@
 def test_import_flag_types(make_input_bundle, get_contract):
     lib1 = """
-import lib2 
+import lib2
 
 flag Roles:
     ADMIN

diff --git a/vyper/codegen/expr.py b/vyper/codegen/expr.py
@@ -207,12 +207,7 @@ def parse_Attribute(self):
         typ = self.expr._metadata["type"]
 
         # (lib1).MyFlag.foo
-        if isinstance(typ, FlagT) and (
-                (isinstance(self.expr.value, vy_ast.Name)
-                 and typ.name == self.expr.value.id)
-                or (isinstance(self.expr.value, vy_ast.Attribute)
-                    and typ.name == self.expr.value.attr)
-        ):
+        if isinstance(typ, FlagT) and is_type_t(self.expr.value._metadata["type"], FlagT):
             # 0, 1, 2, .. 255
             flag_id = typ._flag_members[self.expr.attr]
             value = 2**flag_id  # 0 => 0001, 1 => 0010, 2 => 0100, etc.