docs: Support hidden globals

Summary:
Some globals we don't want to have in docs, specifically `starlark_rust_internals`. So make an option on namespaces and globals to have entries that don't show up in documentation.

I could've added this to more stuff, but felt unnecessary until we actually used it

Reviewed By: stepancheg

Differential Revision: D64858309

fbshipit-source-id: 14eb6c215f65aebb99571d60dd20f50635cae7b4

starlark/src/environment/globals.rs

@@ -36,6 +36,7 @@ pub use crate::stdlib::LibraryExtension;
 use crate::typing::Ty;
 use crate::values::function::NativeFunc;
 use crate::values::function::SpecialBuiltinFunction;
+use crate::values::namespace::value::MaybeDocHiddenValue;
 use crate::values::namespace::FrozenNamespace;
 use crate::values::types::function::NativeFunction;
 use crate::values::AllocFrozenValue;
@@ -49,10 +50,12 @@ use crate::values::Value;
 #[derive(Clone, Dupe, Debug, Allocative)]
 pub struct Globals(Arc<GlobalsData>);
 
+type GlobalValue = MaybeDocHiddenValue<'static, FrozenValue>;
+
 #[derive(Debug, Allocative)]
 struct GlobalsData {
     heap: FrozenHeapRef,
-    variables: SymbolMap<FrozenValue>,
+    variables: SymbolMap<GlobalValue>,
     variable_names: Vec<FrozenStringValue>,
     docstring: Option<String>,
 }
@@ -63,9 +66,9 @@ pub struct GlobalsBuilder {
     // The heap everything is allocated in
     heap: FrozenHeap,
     // Normal top-level variables, e.g. True/hash
-    variables: SymbolMap<FrozenValue>,
+    variables: SymbolMap<GlobalValue>,
     // The list of struct fields, pushed to the end
-    namespace_fields: Vec<SmallMap<FrozenStringValue, FrozenValue>>,
+    namespace_fields: Vec<SmallMap<FrozenStringValue, GlobalValue>>,
     /// The raw docstring for this module
     ///
     /// FIXME(JakobDegen): This should probably be removed. Having a docstring on a `GlobalsBuilder`
@@ -116,7 +119,7 @@ impl Globals {
     /// This function is only safe if you first call `heap` and keep a reference to it.
     /// Therefore, don't expose it on the public API.
     pub(crate) fn get_frozen(&self, name: &str) -> Option<FrozenValue> {
-        self.0.variables.get_str(name).copied()
+        self.0.variables.get_str(name).map(|x| x.value)
     }
 
     /// Get all the names defined in this environment.
@@ -127,7 +130,7 @@ impl Globals {
     /// Iterate over all the items in this environment.
     /// Note returned values are owned by this globals.
     pub fn iter(&self) -> impl Iterator<Item = (&str, FrozenValue)> {
-        self.0.variables.iter().map(|(n, v)| (n.as_str(), *v))
+        self.0.variables.iter().map(|(n, v)| (n.as_str(), v.value))
     }
 
     pub(crate) fn heap(&self) -> &FrozenHeapRef {
@@ -139,7 +142,7 @@ impl Globals {
         self.0
             .variables
             .iter()
-            .map(|(name, val)| val.to_value().describe(name.as_str()))
+            .map(|(name, val)| val.value.to_value().describe(name.as_str()))
             .join("\n")
     }
 
@@ -152,7 +155,11 @@ impl Globals {
     pub fn documentation(&self) -> DocModule {
         let (docs, members) = common_documentation(
             &self.0.docstring,
-            self.0.variables.iter().map(|(n, v)| (n.as_str(), *v)),
+            self.0
+                .variables
+                .iter()
+                .filter(|(_, v)| !v.doc_hidden)
+                .map(|(n, v)| (n.as_str(), v.value)),
         );
         DocModule {
             docs,
@@ -197,10 +204,28 @@ impl GlobalsBuilder {
     /// Add a nested namespace to the builder. If `f` adds the definition `foo`,
     /// it will end up on a namespace `name`, accessible as `name.foo`.
     pub fn namespace(&mut self, name: &str, f: impl FnOnce(&mut GlobalsBuilder)) {
+        self.namespace_inner(name, false, f)
+    }
+
+    /// Same as `namespace`, but this value will not show up in generated documentation.
+    pub fn namespace_no_docs(&mut self, name: &str, f: impl FnOnce(&mut GlobalsBuilder)) {
+        self.namespace_inner(name, true, f)
+    }
+
+    fn namespace_inner(
+        &mut self,
+        name: &str,
+        doc_hidden: bool,
+        f: impl FnOnce(&mut GlobalsBuilder),
+    ) {
         self.namespace_fields.push(SmallMap::new());
         f(self);
         let fields = self.namespace_fields.pop().unwrap();
-        self.set(name, self.heap.alloc(FrozenNamespace::new(fields)));
+        self.set_inner(
+            name,
+            self.heap.alloc(FrozenNamespace::new(fields)),
+            doc_hidden,
+        );
     }
 
     /// A fluent API for modifying [`GlobalsBuilder`] and returning the result.
@@ -234,6 +259,15 @@ impl GlobalsBuilder {
     /// Set a value in the [`GlobalsBuilder`].
     pub fn set<'v, V: AllocFrozenValue>(&'v mut self, name: &str, value: V) {
         let value = value.alloc_frozen_value(&self.heap);
+        self.set_inner(name, value, false)
+    }
+
+    fn set_inner<'v>(&'v mut self, name: &str, value: FrozenValue, doc_hidden: bool) {
+        let value = MaybeDocHiddenValue {
+            value,
+            doc_hidden,
+            phantom: Default::default(),
+        };
         match self.namespace_fields.last_mut() {
             None => {
                 // TODO(nga): do not quietly ignore redefinitions.
@@ -326,15 +360,15 @@ impl GlobalsStatic {
                 .join(", ")
         );
 
-        *globals.0.variables.values().next().unwrap()
+        globals.0.variables.values().next().unwrap().value
     }
 
     /// Move all the globals in this [`GlobalsBuilder`] into a new one. All variables will
     /// only be allocated once (ensuring things like function comparison works properly).
     pub fn populate(&'static self, x: impl FnOnce(&mut GlobalsBuilder), out: &mut GlobalsBuilder) {
         let globals = self.globals(x);
         for (name, value) in globals.0.variables.iter() {
-            out.set(name.as_str(), *value)
+            out.set_inner(name.as_str(), value.value, value.doc_hidden)
         }
         out.docstring = globals.0.docstring.clone();
     }
@@ -357,12 +391,40 @@ pub(crate) fn common_documentation<'a>(
 
 #[cfg(test)]
 mod tests {
+    use starlark_derive::starlark_module;
+
     use super::*;
+    use crate as starlark;
 
     #[test]
     fn test_send_sync()
     where
         Globals: Send + Sync,
     {
     }
+
+    #[starlark_module]
+    fn register_foo(builder: &mut GlobalsBuilder) {
+        fn foo() -> anyhow::Result<i32> {
+            Ok(1)
+        }
+    }
+
+    #[test]
+    fn test_doc_hidden() {
+        let mut globals = GlobalsBuilder::new();
+        globals.namespace_no_docs("ns_hidden", |_| {});
+        globals.namespace("ns", |globals| {
+            globals.namespace_no_docs("nested_ns_hidden", |_| {});
+            globals.set("x", FrozenValue::new_none());
+        });
+        let docs = globals.build().documentation();
+
+        let (k, v) = docs.members.into_iter().exactly_one().ok().unwrap();
+        assert_eq!(&k, "ns");
+        let DocItem::Module(docs) = v else {
+            unreachable!()
+        };
+        assert_eq!(&docs.members.into_keys().exactly_one().ok().unwrap(), "x");
+    }
 }

starlark/src/stdlib/internal.rs

@@ -34,7 +34,7 @@ fn starlark_rust_internal_members(globals: &mut GlobalsBuilder) {
 }
 
 pub(crate) fn register_internal(globals: &mut GlobalsBuilder) {
-    globals.namespace("starlark_rust_internal", |s| {
+    globals.namespace_no_docs("starlark_rust_internal", |s| {
         starlark_rust_internal_members(s)
     });
 }

starlark/src/values/types/namespace.rs

@@ -19,7 +19,7 @@
 
 pub(crate) mod globals;
 mod typing;
-mod value;
+pub(crate) mod value;
 
 pub use value::FrozenNamespace;
 pub use value::Namespace;

starlark/src/values/types/namespace/globals.rs

@@ -21,6 +21,7 @@ use crate as starlark;
 use crate::environment::GlobalsBuilder;
 use crate::eval::Arguments;
 use crate::values::namespace::typing::TyNamespaceFunction;
+use crate::values::namespace::value::MaybeDocHiddenValue;
 use crate::values::namespace::FrozenNamespace;
 use crate::values::namespace::Namespace;
 use crate::values::Heap;
@@ -34,6 +35,20 @@ pub fn register_namespace(builder: &mut GlobalsBuilder) {
     fn namespace<'v>(args: &Arguments<'v, '_>, heap: &'v Heap) -> starlark::Result<Namespace<'v>> {
         args.no_positional_args(heap)?;
 
-        Ok(Namespace::new(args.names_map()?))
+        Ok(Namespace::new(
+            args.names_map()?
+                .into_iter()
+                .map(|(k, v)| {
+                    (
+                        k,
+                        MaybeDocHiddenValue {
+                            value: v,
+                            doc_hidden: false,
+                            phantom: Default::default(),
+                        },
+                    )
+                })
+                .collect(),
+        ))
     }
 }

starlark/src/values/types/namespace/value.rs

@@ -17,6 +17,7 @@
 
 use std::fmt;
 use std::fmt::Display;
+use std::marker::PhantomData;
 
 use allocative::Allocative;
 use display_container::fmt_keyed_container;
@@ -43,20 +44,28 @@ use crate::values::StringValueLike;
 use crate::values::Value;
 use crate::values::ValueLike;
 
+#[derive(Clone, Coerce, Debug, Trace, Freeze, Allocative)]
+#[repr(C)]
+pub(crate) struct MaybeDocHiddenValue<'v, V: ValueLike<'v>> {
+    pub(crate) value: V,
+    pub(crate) doc_hidden: bool,
+    pub(crate) phantom: PhantomData<&'v ()>,
+}
+
 /// The return value of `namespace()`
 #[derive(Clone, Debug, Trace, Freeze, ProvidesStaticType, Allocative)]
 #[repr(C)]
 pub struct NamespaceGen<'v, V: ValueLike<'v>> {
-    fields: SmallMap<V::String, V>,
+    fields: SmallMap<V::String, MaybeDocHiddenValue<'v, V>>,
 }
 
 impl<'v, V: ValueLike<'v>> NamespaceGen<'v, V> {
-    pub fn new(fields: SmallMap<V::String, V>) -> Self {
+    pub(crate) fn new(fields: SmallMap<V::String, MaybeDocHiddenValue<'v, V>>) -> Self {
         Self { fields }
     }
 
     pub fn get(&self, key: &str) -> Option<V> {
-        self.fields.get_hashed(Hashed::new(key)).copied()
+        self.fields.get_hashed(Hashed::new(key)).map(|v| v.value)
     }
 }
 
@@ -71,7 +80,7 @@ impl<'v, V: ValueLike<'v>> Display for NamespaceGen<'v, V> {
             "namespace(",
             ")",
             "=",
-            self.fields.iter().map(|(k, v)| (k.as_str(), v)),
+            self.fields.iter().map(|(k, v)| (k.as_str(), v.value)),
         )
     }
 }
@@ -90,7 +99,9 @@ where
     }
 
     fn get_attr_hashed(&self, attribute: Hashed<&str>, _heap: &'v Heap) -> Option<Value<'v>> {
-        self.fields.get_hashed(attribute).map(|v| v.to_value())
+        self.fields
+            .get_hashed(attribute)
+            .map(|v| v.value.to_value())
     }
 
     fn dir_attr(&self) -> Vec<String> {
@@ -103,7 +114,8 @@ where
             members: self
                 .fields
                 .iter()
-                .map(|(k, v)| (k.as_str().to_owned(), v.to_value().documentation()))
+                .filter(|(_, v)| !v.doc_hidden)
+                .map(|(k, v)| (k.as_str().to_owned(), v.value.to_value().documentation()))
                 .collect(),
         })
     }
@@ -120,7 +132,12 @@ where
             fields: self
                 .fields
                 .iter()
-                .map(|(name, value)| (ArcStr::from(name.as_str()), Ty::of_value(value.to_value())))
+                .map(|(name, value)| {
+                    (
+                        ArcStr::from(name.as_str()),
+                        Ty::of_value(value.value.to_value()),
+                    )
+                })
                 .collect(),
             extra: false,
         }))
@@ -132,7 +149,7 @@ impl<'v, V: ValueLike<'v>> Serialize for NamespaceGen<'v, V> {
     where
         S: serde::Serializer,
     {
-        serializer.collect_map(self.fields.iter())
+        serializer.collect_map(self.fields.iter().map(|(k, v)| (k, v.value)))
     }
 }