compose visibility modifiers with enum class and protocol

Author: KotlinIsland

crates/by_transforms/src/transforms/enums.rs

@@ -193,6 +193,19 @@ fn has_marker(class: &StmtClassDef, source: &str, marker: &str) -> bool {
     })
 }
 
+/// The visibility keyword prefix (`"private "` / `"export "`) to emit ahead of
+/// the enum's synthesized `class` line, or `""` for the default. `public` is
+/// normalized to the `export` marker by the parser.
+fn enum_visibility_prefix(class: &StmtClassDef, source: &str) -> &'static str {
+    if has_marker(class, source, "private") {
+        "private "
+    } else if has_marker(class, source, "export") {
+        "export "
+    } else {
+        ""
+    }
+}
+
 fn check_no_nested_enums(body: &[Stmt], source: &str, nested: bool, errors: &mut Vec<String>) {
     for stmt in body {
         match stmt {
@@ -302,12 +315,20 @@ fn emit_enum(
     let name = class.name.as_str();
     let is_generic = class.type_params.is_some();
 
+    // a `private`/`export` modifier on the enum (`private enum class E:`) rides
+    // through as a keyword prefix on the synthesized base class: phase-1's
+    // `modifiers` pass then renames every module-level reference to the private
+    // name (variant bases `class _E_V(E)`, attachments `E.V`) or registers the
+    // export in `__all__` — the enum owns no visibility logic of its own.
+    // each emitter derives the prefix from the class it already has
+
     // `is_all_unit_enum` is the shared predicate with ty (it models such enums
     // as idiomatic `Enum`s): all variants unit AND no assignment members —
     // python's `Enum` would turn a class-body constant into a *member*,
     // silently diverging from the constant semantics the checker models
     if class.is_all_unit_enum() && !is_generic {
-        emit_plain_enum(out, source, name, &variants, &members, imports);
+        let vis = enum_visibility_prefix(class, source);
+        emit_plain_enum(out, source, name, vis, &variants, &members, imports);
     } else {
         emit_sealed_hierarchy(
             out,
@@ -331,14 +352,15 @@ fn emit_plain_enum(
     out: &mut Out,
     source: &str,
     name: &str,
+    vis: &str,
     variants: &[Variant],
     members: &[&Stmt],
     imports: &mut ImportSet,
 ) {
     imports.add("enum", "Enum");
     imports.add("enum", "auto");
 
-    out.push_gen(&format!("class {name}(Enum):\n"));
+    out.push_gen(&format!("{vis}class {name}(Enum):\n"));
     for variant in variants {
         out.push_gen(&format!("{}{} = auto()\n", variant.indent, variant.name));
     }
@@ -369,6 +391,9 @@ fn emit_sealed_hierarchy(
     imports: &mut ImportSet,
     min_version: PythonVersion,
 ) {
+    // visibility prefix derived from the class (see `lower`): `private`/`export`
+    // ride through to phase-1's `modifiers` pass on the synthesized base line
+    let vis = enum_visibility_prefix(class, source);
     // payload variants lower to `@final` frozen dataclasses; unit variants are
     // plain singleton values needing neither import
     if variants.iter().any(|v| v.kind != VariantKind::Unit) {
@@ -385,7 +410,7 @@ fn emit_sealed_hierarchy(
         .map(|tp| slice(source, tp.range()).to_string())
         .unwrap_or_default();
 
-    out.push_gen(&format!("class {name}{params}:\n"));
+    out.push_gen(&format!("{vis}class {name}{params}:\n"));
     // ordinary members (methods, classmethods, constants) — copied verbatim,
     // already indented under the enum in the source. they may refer to variants
     // (`A.Foo`) freely: the references resolve lazily at call time, by which
@@ -620,6 +645,54 @@ mod tests {
         );
     }
 
+    #[test]
+    fn export_enum_registers_in_all() {
+        // `export enum class` rides through as an `export class` prefix on the
+        // synthesized enum, so phase-1's modifiers pass registers it in `__all__`
+        check(
+            indoc! {"
+                export enum class Color:
+                    case Red, Green
+            "},
+            indoc! {"
+                from __future__ import annotations
+                from enum import Enum, auto
+                class Color(Enum):
+                    Red = auto()
+                    Green = auto()
+                __all__ = [\"Color\"]
+            "},
+        );
+    }
+
+    #[test]
+    fn private_enum_renames_whole_hierarchy() {
+        // `private enum class` becomes a `private class` prefix; modifiers renames
+        // the declaration *and* every module-level reference (the variant bases and
+        // attachments the enum synthesized), so the private name stays consistent
+        check(
+            indoc! {"
+                private enum class Shape:
+                    case Circle(radius: int)
+            "},
+            indoc! {"
+                from __future__ import annotations
+                from dataclasses import dataclass
+                from typing import final
+                class _Shape:
+                    pass
+
+                @final
+                @dataclass(frozen=True, slots=True)
+                class _Shape_Circle(_Shape):
+                    radius: int
+                _Shape_Circle.__name__ = \"Circle\"
+                _Shape_Circle.__qualname__ = \"Shape.Circle\"
+                _Shape.Circle = _Shape_Circle
+            "},
+        );
+    }
+
     #[test]
     fn all_unit_with_constant_avoids_enum_form() {
         // python's `Enum` turns every class-body assignment into a *member* —

crates/by_transforms/src/transforms/modifiers.rs

@@ -823,6 +823,31 @@ mod tests {
         );
     }
 
+    #[test]
+    fn private_protocol() {
+        // a visibility modifier composes with the `protocol` introducer: both the
+        // `private` rename and the `protocol_class` rewrite lower by disjoint edits
+        check(
+            "private protocol Foo: ...\n",
+            indoc! {"
+                from typing import Protocol
+                class _Foo(Protocol): ...
+            "},
+        );
+    }
+
+    #[test]
+    fn export_protocol() {
+        check(
+            "export protocol Foo: ...\n",
+            indoc! {"
+                from typing import Protocol
+                class Foo(Protocol): ...
+                __all__ = [\"Foo\"]
+            "},
+        );
+    }
+
     #[test]
     fn let_decl() {
         check(

crates/ruff_python_parser/src/parser/statement.rs

@@ -247,7 +247,7 @@ impl<'src> Parser<'src> {
             self.error_if_not_basedpython(
                 "`protocol` class syntax is not valid in .py files".to_string(),
             );
-            return Some(self.parse_protocol_def(start));
+            return Some(self.parse_protocol_def(start, DecoratorList::new()));
         }
         // `enum class E:` / `enum class E[T]:` — a "based enum" (an algebraic
         // sum type when its body has payload variants, an idiomatic `Enum` when
@@ -258,7 +258,7 @@ impl<'src> Parser<'src> {
             self.error_if_not_basedpython(
                 "`enum class` declarations are not valid in .py files".to_string(),
             );
-            return Some(self.parse_enum_def(start));
+            return Some(self.parse_enum_def(start, DecoratorList::new()));
         }
         // a bare `enum E:` (no `class`) is not valid — based enums are written
         // `enum class E:`. report it but still parse the body for recovery
@@ -269,7 +269,7 @@ impl<'src> Parser<'src> {
                 ),
                 self.current_token_range(),
             );
-            return Some(self.parse_enum_def(start));
+            return Some(self.parse_enum_def(start, DecoratorList::new()));
         }
         if kw == "decorator" && self.peek() == TokenKind::Def {
             self.error_if_not_basedpython(
@@ -358,6 +358,18 @@ impl<'src> Parser<'src> {
                     } else {
                         self.peek_nth(idx).0
                     };
+                    // an introducer keyword (`enum class`, `protocol`) after a
+                    // modifier chain — `private enum class E:`, `export protocol P:`.
+                    // dispatch through the modifier path so the chain is carried
+                    // as decorators on the introduced class
+                    if (text == "protocol" && following == TokenKind::Name)
+                        || (text == "enum" && following == TokenKind::Class)
+                    {
+                        self.error_if_not_basedpython(format!(
+                            "`{kw}` is a basedpython modifier and is not valid in .py files"
+                        ));
+                        return Some(self.parse_with_modifier(start, DecoratorList::new()));
+                    }
                     return match following {
                         TokenKind::Equal if idx > 0 => {
                             self.error_if_not_basedpython(format!(
@@ -524,7 +536,8 @@ impl<'src> Parser<'src> {
             if self.at(TokenKind::Def) || self.at(TokenKind::Class) || self.at(TokenKind::Async) {
                 break;
             }
-            // another modifier keyword follows — keep looping
+            // another modifier keyword follows — keep looping. an `enum class` /
+            // `protocol` introducer also ends the chain (handled after the loop)
             if self.at(TokenKind::Name) {
                 let kw = self.src_text(self.current_token_range());
                 // any modifier keyword may follow another, in any order — reuse
@@ -548,6 +561,20 @@ impl<'src> Parser<'src> {
             })
         } else if self.at(TokenKind::Def) {
             Stmt::FunctionDef(self.parse_function_definition(decorators, start))
+        } else if self.at(TokenKind::Name)
+            && self.src_text(self.current_token_range()) == "protocol"
+            && self.peek() == TokenKind::Name
+        {
+            // `private protocol P:` — the modifier decorators ride alongside the
+            // synthetic `protocol_class` marker; both lower by disjoint edits
+            self.parse_protocol_def(start, decorators)
+        } else if self.at(TokenKind::Name)
+            && self.src_text(self.current_token_range()) == "enum"
+            && self.peek() == TokenKind::Class
+        {
+            // `private enum class E:` — the enum lowering reads the visibility
+            // markers to prefix its synthesized `class` line
+            self.parse_enum_def(start, decorators)
         } else {
             Stmt::ClassDef(self.parse_class_definition(decorators, start))
         }
@@ -853,13 +880,16 @@ impl<'src> Parser<'src> {
     /// requiring an explicit `class` keyword. Produces a `ClassDef` with a synthetic
     /// `Decorator` (name `"protocol_class"`) that the `modifiers` transform rewrites
     /// to `class Foo(Protocol):`.
-    fn parse_protocol_def(&mut self, start: TextSize) -> Stmt {
+    fn parse_protocol_def(&mut self, start: TextSize, mut decorators: DecoratorList) -> Stmt {
         let protocol_start = self.current_token_range().start();
         self.bump(TokenKind::Name); // consume "protocol"
         let class_name_start = self.current_token_range().start();
         let decorator_range = TextRange::new(protocol_start, class_name_start);
 
-        let decorator = ast::Decorator {
+        // the synthetic `protocol_class` marker follows any modifier decorators
+        // (`private protocol P:`); the `modifiers` transform consumes each by a
+        // disjoint range edit, so order between them does not matter
+        decorators.push(ast::Decorator {
             expression: Expr::Name(ast::ExprName {
                 id: Name::new_static("protocol_class"),
                 ctx: ExprContext::Invalid,
@@ -868,7 +898,7 @@ impl<'src> Parser<'src> {
             }),
             range: decorator_range,
             node_index: AtomicNodeIndex::NONE,
-        };
+        });
 
         let name = self.parse_identifier();
         let type_params = self.try_parse_type_params();
@@ -884,7 +914,7 @@ impl<'src> Parser<'src> {
 
         Stmt::ClassDef(ast::StmtClassDef {
             range: self.node_range(start),
-            decorator_list: vec![decorator].into(),
+            decorator_list: decorators,
             name,
             type_params: type_params.map(Box::new),
             arguments,
@@ -903,7 +933,7 @@ impl<'src> Parser<'src> {
     ///
     /// [`ClassDef`]: ast::StmtClassDef
     /// [`AnnAssign`]: ast::StmtAnnAssign
-    fn parse_enum_def(&mut self, start: TextSize) -> Stmt {
+    fn parse_enum_def(&mut self, start: TextSize, mut decorators: DecoratorList) -> Stmt {
         let enum_start = self.current_token_range().start();
         self.bump(TokenKind::Name); // consume "enum"
         // the canonical surface is `enum class E:`; the `class` keyword is part
@@ -913,7 +943,11 @@ impl<'src> Parser<'src> {
         let name_start = self.current_token_range().start();
         let decorator_range = TextRange::new(enum_start, name_start);
 
-        let decorator = ast::Decorator {
+        // the synthetic `enum_def` marker follows any modifier decorators
+        // (`private enum class E:`); the enum lowering re-emits the enum from
+        // scratch and reads the visibility markers to prefix the synthesized
+        // `class` line, so the standard `private`/`export` class path applies
+        decorators.push(ast::Decorator {
             expression: Expr::Name(ast::ExprName {
                 id: Name::new_static("enum_def"),
                 ctx: ExprContext::Invalid,
@@ -922,7 +956,7 @@ impl<'src> Parser<'src> {
             }),
             range: decorator_range,
             node_index: AtomicNodeIndex::NONE,
-        };
+        });
 
         let name = self.parse_identifier();
         let type_params = self.try_parse_type_params();
@@ -952,7 +986,7 @@ impl<'src> Parser<'src> {
 
         Stmt::ClassDef(ast::StmtClassDef {
             range: self.node_range(start),
-            decorator_list: vec![decorator].into(),
+            decorator_list: decorators,
             name,
             type_params: type_params.map(Box::new),
             arguments: None,

crates/ty_python_semantic/resources/mdtest/basedpython_enums.md

@@ -272,6 +272,27 @@ def f(e: WithConst) -> str:
             return "b"
 ```
 
+## visibility modifiers compose with `enum class`
+
+a `private` or `export` modifier may precede `enum class`. `export` registers the enum in the
+module's `__all__`; `private` renames the enum — and every reference the lowering synthesizes (the
+variant subclasses, their attachments) — to an underscore-prefixed module-private name, so the
+public surface stays consistent. the variants are still reached qualified through the declared name.
+
+```by
+export enum class Color:
+    case Red, Green
+
+private enum class Shape:
+    case Circle(radius: float)
+    case Square(side: float)
+
+reveal_type(Color.Red)  # revealed: Color.Red
+c = Shape.Circle(2.0)
+reveal_type(c.radius)  # revealed: float
+reveal_type(Shape.Square(1.0))  # revealed: Square
+```
+
 ## variants require `case`
 
 a bare name in an `enum class` body is a no-op statement, almost certainly a variant missing its

crates/ty_python_semantic/resources/mdtest/basedpython_visibility.md

@@ -0,0 +1,50 @@
+# basedpython: visibility modifiers
+
+`private` and `export`/`public` are transpile-time visibility modifiers: `export`/`public` add the
+symbol to the module's generated `__all__`, and `private` renames it with an underscore prefix (or,
+inside a class body, name-mangles it with `__`). they carry no type-level effect — the decorated
+class or function keeps its ordinary type rather than being erased to `Unknown`.
+
+## a private class keeps its type
+
+```by
+private class Box:
+    value: int = 0
+
+b = Box()
+reveal_type(b)  # revealed: Box
+reveal_type(b.value)  # revealed: int
+```
+
+## an exported function keeps its signature
+
+```by
+export def make(n: int) -> int:
+    return n * 2
+
+reveal_type(make)  # revealed: def make(n: int) -> int
+reveal_type(make(3))  # revealed: int
+```
+
+## `open` is a no-op modifier on the type
+
+`open` marks a class freely subclassable (the default in Python); it only suppresses the closed-by-
+default checks, so the class type is unchanged.
+
+```by
+open class Base:
+    tag: str = "b"
+
+reveal_type(Base().tag)  # revealed: str
+```
+
+## modifiers compose with each other
+
+a chain of modifiers still resolves to the underlying declaration's type.
+
+```by
+private final class Sealed:
+    n: int = 1
+
+reveal_type(Sealed().n)  # revealed: int
+```