From 7818ffdcad86fbe08a31b6f71160f9e6020c0e33 Mon Sep 17 00:00:00 2001
From: oftheforest <usual.salt7109@fastmail.com>
Date: Sat, 27 Apr 2024 11:25:06 -0700
Subject: [PATCH] Allow subenum visibility to be overridden

---
 README.md    | 34 +++++++++++++++++++++++++++++++++
 src/build.rs |  2 +-
 src/enum.rs  | 11 +++++++++--
 src/lib.rs   | 54 ++++++++++++++++++++++++++++++++++++++--------------
 4 files changed, 84 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index e3749c3..13b4540 100644
--- a/README.md
+++ b/README.md
@@ -134,6 +134,40 @@ fn main() {
 }
 ```
 
+## Subenum visibility
+
+By default, a subenum has the same visibility as its parent enum. Use `vis` to
+override this.
+
+```rust
+use subenum::subenum;
+
+#[subenum(Bar(vis = "pub"))]
+pub(crate) enum Foo {
+    #[subenum(Bar)]
+    A(String),
+    B,
+    #[subenum(Bar)]
+    C(u8),
+}
+```
+
+The visibility qualifier for "private" is empty, so use `vis = ""` to make a
+subenum private.
+
+```rust
+use subenum::subenum;
+
+#[subenum(Bar(vis = ""))]
+pub enum Foo {
+    #[subenum(Bar)]
+    A(String),
+    B,
+    #[subenum(Bar)]
+    C(u8),
+}
+```
+
 
 # Limitations
 
diff --git a/src/build.rs b/src/build.rs
index 68e79c7..529601e 100644
--- a/src/build.rs
+++ b/src/build.rs
@@ -136,7 +136,7 @@ impl Enum {
             .iter()
             .map(|&derive| self.build_inherited_derive(parent, derive, &self.variants));
 
-        let vis = &parent.vis;
+        let vis = self.visibility.as_ref().unwrap_or(&parent.vis);
 
         let (_child_impl, child_ty, _child_where) = child_generics.split_for_impl();
 
diff --git a/src/enum.rs b/src/enum.rs
index 2a3138d..36c8bb2 100644
--- a/src/enum.rs
+++ b/src/enum.rs
@@ -3,7 +3,7 @@ use alloc::{
     vec::Vec,
 };
 use proc_macro2::TokenStream;
-use syn::{punctuated::Punctuated, Generics, Ident, Token, TypeParamBound, Variant};
+use syn::{punctuated::Punctuated, Generics, Ident, Token, TypeParamBound, Variant, Visibility};
 
 use crate::{extractor::Extractor, iter::BoxedIter, param::Param, Derive};
 
@@ -14,10 +14,16 @@ pub struct Enum {
     pub attributes: Vec<TokenStream>,
     pub derives: Vec<Derive>,
     pub generics: Generics,
+    pub visibility: Option<Visibility>,
 }
 
 impl Enum {
-    pub fn new(ident: Ident, attributes: Vec<TokenStream>, derives: Vec<Derive>) -> Self {
+    pub fn new(
+        ident: Ident,
+        visibility: Option<Visibility>,
+        attributes: Vec<TokenStream>,
+        derives: Vec<Derive>,
+    ) -> Self {
         Enum {
             ident,
             variants: Punctuated::new(),
@@ -30,6 +36,7 @@ impl Enum {
                 gt_token: Some(syn::token::Gt::default()),
                 where_clause: None,
             },
+            visibility,
         }
     }
 
diff --git a/src/lib.rs b/src/lib.rs
index 735174e..71d6e64 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -19,8 +19,8 @@ use proc_macro2::Ident;
 use quote::quote;
 use r#enum::Enum;
 use syn::{
-    parse_macro_input, Attribute, AttributeArgs, DeriveInput, Field, Meta, MetaList, MetaNameValue,
-    NestedMeta, Type,
+    parse_macro_input, Attribute, AttributeArgs, DeriveInput, Field, Lit, Meta, MetaList,
+    MetaNameValue, NestedMeta, Type,
 };
 
 const SUBENUM: &str = "subenum";
@@ -77,28 +77,54 @@ fn build_enum_map(args: AttributeArgs, derives: &[Derive]) -> BTreeMap<Ident, En
             NestedMeta::Lit(_) => panic!("{}", ERR),
         })
         .map(|meta| match meta {
-            Meta::Path(path) => (path.get_ident().expect(ERR).to_owned(), Vec::new()),
-            Meta::List(MetaList { path, nested, .. }) => (
-                path.get_ident().expect(ERR).to_owned(),
-                nested
+            Meta::Path(path) => (path.get_ident().expect(ERR).to_owned(), Vec::new(), None),
+            Meta::List(MetaList { path, nested, .. }) => {
+                let mut vis = None;
+                let attrs = nested
                     .into_iter()
                     .map(|nested| match nested {
                         NestedMeta::Meta(meta) => meta,
                         NestedMeta::Lit(_) => panic!("{}", ERR),
                     })
-                    .map(|meta| match meta {
-                        Meta::Path(path) => quote! { #path },
-                        Meta::List(MetaList { path, nested, .. }) => quote! { #path(#nested) },
-                        Meta::NameValue(MetaNameValue { path, lit, .. }) => quote! { #path = #lit },
+                    .filter_map(|meta| match meta {
+                        Meta::Path(path) => Some(quote! { #path }),
+                        Meta::List(MetaList { path, nested, .. }) => {
+                            Some(quote! { #path(#nested) })
+                        }
+                        Meta::NameValue(MetaNameValue { path, lit, .. }) => {
+                            if path.is_ident("vis") {
+                                if vis.is_some() {
+                                    panic!("subenum has conflicting visibility qualifiers");
+                                }
+                                let string = match lit {
+                                    Lit::Str(lit_str) => lit_str,
+                                    _ => panic!(r#"subenum visibility must be passed as a string, e.g. `#[subenum(EnumA, EnumB(vis = "pub")]`"#),
+                                };
+                                let new_vis = string
+                                    .parse()
+                                    .unwrap_or_else(|_| panic!(
+                                        "subenum vis is not a valid visibility qualifier",
+                                    ));
+                                vis = Some(new_vis);
+                                None
+                            } else {
+                                Some(quote! { #path = #lit })
+                            }
+                        }
                     })
-                    .collect::<Vec<proc_macro2::TokenStream>>(),
-            ),
+                    .collect::<Vec<proc_macro2::TokenStream>>();
+                (
+                    path.get_ident().expect(ERR).to_owned(),
+                    attrs,
+                    vis,
+                )
+            },
             _ => panic!("{}", ERR),
         })
-        .map(|(ident, attrs)| {
+        .map(|(ident, attrs, vis)| {
             (
                 ident.clone(),
-                Enum::new(ident.clone(), attrs, derives.to_owned()),
+                Enum::new(ident.clone(), vis, attrs, derives.to_owned()),
             )
         })
         .collect()