Data: Allow undefined return from withSelect mapSelectToProps

data/index.js

@@ -218,14 +218,18 @@ export const withSelect = ( mapStateToProps ) => ( WrappedComponent ) => {
 				return;
 			}
 
-			const newState = mapStateToProps( select, props );
-			if ( ! isEqualShallow( newState, this.state ) ) {
-				this.setState( newState );
+			const { mergeProps } = this.state;
+			const nextMergeProps = mapStateToProps( select, props ) || {};
+
+			if ( ! isEqualShallow( nextMergeProps, mergeProps ) ) {
+				this.setState( {
+					mergeProps: nextMergeProps,
+				} );
 			}
 		}
 
 		render() {
-			return <WrappedComponent { ...this.props } { ...this.state } />;
+			return <WrappedComponent { ...this.props } { ...this.state.mergeProps } />;
 		}
 	}
 

data/test/index.js

@@ -273,6 +273,54 @@ describe( 'withSelect', () => {
 
 			expect( mapSelectToProps ).toHaveBeenCalledTimes( 1 );
 		} );
+
+		itWithExtraAssertions( 'omits props which are not returned on subsequent mappings', () => {
+			registerReducer( 'demo', ( state = 'OK' ) => state );
+			registerSelectors( 'demo', {
+				getValue: ( state ) => state,
+			} );
+
+			const Component = withSelectImplementation( ( _select, ownProps ) => {
+				return {
+					[ ownProps.propName ]: _select( 'demo' ).getValue(),
+				};
+			} )( () => <div /> );
+
+			wrapper = mount( <Component propName="foo" /> );
+
+			expect( wrapper.childAt( 0 ).props() ).toEqual( { foo: 'OK', propName: 'foo' } );
+
+			wrapper.setProps( { propName: 'bar' } );
+
+			expect( wrapper.childAt( 0 ).props() ).toEqual( { bar: 'OK', propName: 'bar' } );
+		} );
+
+		itWithExtraAssertions( 'allows undefined return from mapSelectToProps', () => {
+			registerReducer( 'demo', ( state = 'OK' ) => state );
+			registerSelectors( 'demo', {
+				getValue: ( state ) => state,
+			} );
+
+			const Component = withSelectImplementation( ( _select, ownProps ) => {
+				if ( ownProps.pass ) {
+					return {
+						count: _select( 'demo' ).getValue(),
+					};
+				}
+			} )( ( props ) => <div>{ props.count || 'Unknown' }</div> );
+
+			wrapper = mount( <Component pass={ false } /> );
+
+			expect( wrapper.childAt( 0 ).text() ).toBe( 'Unknown' );
+
+			wrapper.setProps( { pass: true } );
+
+			expect( wrapper.childAt( 0 ).text() ).toBe( 'OK' );
+
+			wrapper.setProps( { pass: false } );
+
+			expect( wrapper.childAt( 0 ).text() ).toBe( 'Unknown' );
+		} );
 	}
 
 	cases( withSelect );