feat: add layout and screenLayout props for screens

example/src/Screens/Layouts.tsx

@@ -0,0 +1,150 @@
+import { Button } from '@react-navigation/elements';
+import type { ParamListBase } from '@react-navigation/native';
+import {
+  createStackNavigator,
+  type StackNavigationOptions,
+  type StackScreenProps,
+} from '@react-navigation/stack';
+import * as React from 'react';
+import { ScrollView, StyleSheet, Text, View } from 'react-native';
+
+export type SimpleStackParams = {
+  SuspenseDemo: { author: string } | undefined;
+};
+
+let cached: number | undefined;
+
+const createPromise = () =>
+  new Promise<number>((resolve) => {
+    setTimeout(() => resolve(42), 1000);
+  }).then((result) => {
+    cached = result;
+    return result;
+  });
+
+const SuspenseDemoScreen = ({
+  navigation,
+}: StackScreenProps<SimpleStackParams, 'SuspenseDemo'>) => {
+  const [promise, setPromise] = React.useState(createPromise);
+  const [error, setError] = React.useState<Error | null>(null);
+
+  // Naive implementation for suspense intended for demo purposes
+  // We need to suspend when there's no cached value by throwing a promise
+  if (cached == null) {
+    throw promise;
+  }
+
+  if (error) {
+    throw error;
+  }
+
+  return (
+    <ScrollView>
+      <View style={styles.buttons}>
+        <Button
+          variant="filled"
+          onPress={() => {
+            cached = undefined;
+            setPromise(createPromise());
+          }}
+        >
+          Suspend
+        </Button>
+        <Button
+          variant="tinted"
+          onPress={() => {
+            setError(new Error('Something went wrong'));
+          }}
+        >
+          Crash
+        </Button>
+        <Button variant="tinted" onPress={() => navigation.goBack()}>
+          Go back
+        </Button>
+      </View>
+    </ScrollView>
+  );
+};
+
+class ErrorBoundary extends React.Component<
+  { children: React.ReactNode },
+  { hasError: boolean }
+> {
+  static getDerivedStateFromError() {
+    return { hasError: true };
+  }
+
+  state = { hasError: false };
+
+  render() {
+    if (this.state.hasError) {
+      return (
+        <View style={styles.fallback}>
+          <Text style={styles.text}>Something went wrong</Text>
+          <Button onPress={() => this.setState({ hasError: false })}>
+            Try again
+          </Button>
+        </View>
+      );
+    }
+
+    return this.props.children;
+  }
+}
+
+const Stack = createStackNavigator<SimpleStackParams>();
+
+export function LayoutsStack({
+  navigation,
+  screenOptions,
+}: StackScreenProps<ParamListBase> & {
+  screenOptions?: StackNavigationOptions;
+}) {
+  React.useLayoutEffect(() => {
+    navigation.setOptions({
+      headerShown: false,
+    });
+  }, [navigation]);
+
+  return (
+    <Stack.Navigator screenOptions={screenOptions}>
+      <Stack.Screen
+        name="SuspenseDemo"
+        component={SuspenseDemoScreen}
+        options={{ title: 'Suspense & ErrorBoundary' }}
+        layout={({ children }) => (
+          <ErrorBoundary>
+            <React.Suspense
+              fallback={
+                <View style={styles.fallback}>
+                  <Text style={styles.text}>Loading…</Text>
+                </View>
+              }
+            >
+              {children}
+            </React.Suspense>
+          </ErrorBoundary>
+        )}
+      />
+    </Stack.Navigator>
+  );
+}
+
+const styles = StyleSheet.create({
+  buttons: {
+    flexDirection: 'row',
+    flexWrap: 'wrap',
+    gap: 12,
+    margin: 12,
+  },
+  fallback: {
+    flex: 1,
+    alignItems: 'center',
+    justifyContent: 'center',
+  },
+
+  text: {
+    textAlign: 'center',
+    margin: 12,
+  },
+});

example/src/screens.tsx

@@ -5,6 +5,7 @@ import { BottomTabs } from './Screens/BottomTabs';
 import { CustomLayout } from './Screens/CustomLayout';
 import { DrawerView } from './Screens/DrawerView';
 import { DynamicTabs } from './Screens/DynamicTabs';
+import { LayoutsStack } from './Screens/Layouts';
 import { LinkComponent } from './Screens/LinkComponent';
 import { LinkingScreen } from './Screens/LinkingScreen';
 import { MasterDetail } from './Screens/MasterDetail';
@@ -85,6 +86,10 @@ export const SCREENS = {
     title: 'Auth Flow',
     component: AuthFlow,
   },
+  Layouts: {
+    title: 'Custom Layout',
+    component: LayoutsStack,
+  },
   StackPreventRemove: {
     title: 'Prevent removing screen in Stack',
     component: StackPreventRemove,

packages/core/src/__tests__/useDescriptors.test.tsx

@@ -7,6 +7,7 @@ import { act, render } from '@testing-library/react-native';
 import * as React from 'react';
 
 import { BaseNavigationContainer } from '../BaseNavigationContainer';
+import { Group } from '../Group';
 import { Screen } from '../Screen';
 import { useNavigationBuilder } from '../useNavigationBuilder';
 import {
@@ -383,6 +384,128 @@ it('updates options with setOptions', () => {
       `);
 });
 
+it('renders layout defined for the screen', () => {
+  const TestNavigator = (props: any) => {
+    const { state, descriptors } = useNavigationBuilder<
+      NavigationState,
+      any,
+      any,
+      any,
+      any
+    >(MockRouter, props);
+    const { render } = descriptors[state.routes[state.index].key];
+
+    return render();
+  };
+
+  const TestScreen = () => {
+    return <>Test screen</>;
+  };
+
+  const element = (
+    <BaseNavigationContainer>
+      <TestNavigator
+        screenLayout={({ children }: any) => <main>{children}</main>}
+      >
+        <Group screenLayout={({ children }) => <section>{children}</section>}>
+          <Screen
+            name="foo"
+            component={TestScreen}
+            layout={({ children }) => <div>{children}</div>}
+          />
+          <Screen name="bar" component={React.Fragment} />
+        </Group>
+      </TestNavigator>
+    </BaseNavigationContainer>
+  );
+
+  const root = render(element);
+
+  expect(root).toMatchInlineSnapshot(`
+<div>
+  Test screen
+</div>
+`);
+});
+
+it('renders layout defined for the group', () => {
+  const TestNavigator = (props: any) => {
+    const { state, descriptors } = useNavigationBuilder<
+      NavigationState,
+      any,
+      any,
+      any,
+      any
+    >(MockRouter, props);
+    const { render } = descriptors[state.routes[state.index].key];
+
+    return render();
+  };
+
+  const TestScreen = () => {
+    return <>Test screen</>;
+  };
+
+  const element = (
+    <BaseNavigationContainer>
+      <TestNavigator
+        screenLayout={({ children }: any) => <main>{children}</main>}
+      >
+        <Group screenLayout={({ children }) => <section>{children}</section>}>
+          <Screen name="foo" component={TestScreen} />
+          <Screen name="bar" component={React.Fragment} />
+        </Group>
+      </TestNavigator>
+    </BaseNavigationContainer>
+  );
+
+  const root = render(element);
+
+  expect(root).toMatchInlineSnapshot(`
+<section>
+  Test screen
+</section>
+`);
+});
+
+it('renders layout defined for the navigator', () => {
+  const TestNavigator = (props: any) => {
+    const { state, descriptors } = useNavigationBuilder<
+      NavigationState,
+      any,
+      any,
+      any,
+      any
+    >(MockRouter, props);
+    const { render } = descriptors[state.routes[state.index].key];
+
+    return render();
+  };
+
+  const TestScreen = () => {
+    return <>Test screen</>;
+  };
+
+  const element = (
+    <BaseNavigationContainer>
+      <TestNavigator
+        screenLayout={({ children }: any) => <main>{children}</main>}
+      >
+        <Screen name="foo" component={TestScreen} />
+        <Screen name="bar" component={React.Fragment} />
+      </TestNavigator>
+    </BaseNavigationContainer>
+  );
+
+  const root = render(element);
+
+  expect(root).toMatchInlineSnapshot(`
+<main>
+  Test screen
+</main>
+`);
+});
+
 it("returns correct value for canGoBack when it's not overridden", () => {
   const TestNavigator = (props: any) => {
     const { state, descriptors } = useNavigationBuilder<

packages/core/src/types.tsx

@@ -32,13 +32,15 @@ export type DefaultNavigatorOptions<
    * Optional ID for the navigator. Can be used with `navigation.getParent(id)` to refer to a parent.
    */
   id?: string;
+
   /**
    * Children React Elements to extract the route configuration from.
    * Only `Screen`, `Group` and `React.Fragment` are supported as children.
    */
   children: React.ReactNode;
+
   /**
-   * Layout component for the navigator.
+   * Layout for the navigator.
    * Useful for wrapping with a component with access to navigator's state and options.
    */
   layout?: (props: {
@@ -61,6 +63,7 @@ export type DefaultNavigatorOptions<
     >;
     children: React.ReactNode;
   }) => React.ReactElement;
+
   /**
    * Event listeners for all the screens in the navigator.
    */
@@ -70,6 +73,7 @@ export type DefaultNavigatorOptions<
         route: RouteProp<ParamList>;
         navigation: any;
       }) => ScreenListeners<State, EventMap>);
+
   /**
    * Default options for all screens under this navigator.
    */
@@ -80,6 +84,17 @@ export type DefaultNavigatorOptions<
         navigation: any;
         theme: ReactNavigation.Theme;
       }) => ScreenOptions);
+
+  /**
+   * Layout for all screens under this navigator.
+   */
+  screenLayout?: (props: {
+    route: RouteProp<ParamList, keyof ParamList>;
+    navigation: any;
+    theme: ReactNavigation.Theme;
+    children: React.ReactElement;
+  }) => React.ReactElement;
+
   /**
    A function returning a state, which may be set after modifying the routes name.
    */
@@ -643,6 +658,18 @@ export type RouteConfig<
         navigation: any;
       }) => ScreenListeners<State, EventMap>);
 
+  /**
+   * Layout for this screen.
+   * Useful for wrapping the screen with custom containers.
+   * e.g. for styling, error boundaries, suspense, etc.
+   */
+  layout?: (props: {
+    route: RouteProp<ParamList, keyof ParamList>;
+    navigation: any;
+    theme: ReactNavigation.Theme;
+    children: React.ReactElement;
+  }) => React.ReactElement;
+
   /**
    * Function to return an unique ID for this screen.
    * Receives an object with the route params.
@@ -681,6 +708,18 @@ export type RouteGroupConfig<
         navigation: any;
         theme: ReactNavigation.Theme;
       }) => ScreenOptions);
+
+  /**
+   * Layout for the screens inside the group.
+   * This will override the `screenLayout` of parent group or navigator.
+   */
+  screenLayout?: (props: {
+    route: RouteProp<ParamList, keyof ParamList>;
+    navigation: any;
+    theme: ReactNavigation.Theme;
+    children: React.ReactElement;
+  }) => React.ReactElement;
+
   /**
    * Children React Elements to extract the route configuration from.
    * Only `Screen`, `Group` and `React.Fragment` are supported as children.