Add class-like state behavior option to withState

docs/API.md

@@ -243,13 +243,14 @@ const Post = enhance(({ title, content, author }) =>
 
 ```js
 withState(
-  stateName: string,
-  stateUpdaterName: string,
+  stateName: ?string,
+  stateUpdaterName: ?string,
   initialState: any | (props: Object) => any
 ): HigherOrderComponent
 ```
+Passes two additional props to the base component: a state value, and a function to update that state value.
 
-Passes two additional props to the base component: a state value, and a function to update that state value. The state updater has the following signature:
+If stateName and stateUpdateName are provided then the behavior is as follows. The state updater has the following signature:
 
 ```js
 stateUpdater<T>((prevValue: T) => T, ?callback: Function): void
@@ -273,7 +274,20 @@ The second form accepts a single value, which is used as the new state.
 
 Both forms accept an optional second parameter, a callback function that will be executed once `setState()` is completed and the component is re-rendered.
 
-An initial state value is required. It can be either the state value itself, or a function that returns an initial state given the initial props.
+Only the initialState param is required. It can be either the state value itself, or a function that returns an initial state given the initial props.
+
+If stateName and stateUpdaterName are not provided then they default to `'state'` and `'setState'`. In this form, state and setState behave exactly like setState in a typical react class. So `state` will be an object and `setState()` will accept a new state object to merge with the existing state. In this form `setState()` also still allows either an object or function as a param for updating state:
+
+```js
+const addCounting = compose(
+  withState({ count: 0 }),
+  withHandlers({
+    increment: ({ setState }) => () => setState(({ count }) => count + 1),
+    decrement: ({ setState }) => () =>  setState(({ count }) => count - 1),
+    reset: ({ setState }) => () => setState({ count: 0 })
+  }))
+)
+```
 
 ### `withReducer()`
 

src/packages/recompose/__tests__/withState-test.js

@@ -1,6 +1,6 @@
 import React from 'react'
 import { withState } from '../'
-import { mount } from 'enzyme'
+import { mount, shallow } from 'enzyme'
 import sinon from 'sinon'
 
 test('withState adds a stateful value and a function for updating it', () => {
@@ -68,3 +68,69 @@ test('withState also accepts initialState as function of props', () => {
   updateCounter(n => n * 3)
   expect(component.lastCall.args[0].counter).toBe(3)
 })
+
+test('withState (1 param) adds state and setState props', () => {
+  const enhance = withState({ clicked: 'no' })
+  const Component = enhance(({ state, setState }) =>
+    <button
+      onClick={() => setState({ clicked: 'yes' })}
+    >
+      {state.clicked}{state.foo}
+    </button>
+  )
+
+  const wrapper = shallow(<Component />)
+
+  expect(wrapper.text()).toBe('no')
+  wrapper.find('button').simulate('click')
+  expect(wrapper.text()).toBe('yes')
+})
+
+test('withState (1 param) allows modify and merge in same update', () => {
+  const enhance = withState({ foo: 'empty' })
+  const Component = enhance(({ state, setState }) =>
+    <button
+      onClick={() => setState({ foo: 'foo', bar: 'bar' })}
+    >
+      {state.foo}{state.bar}
+    </button>
+  )
+
+  const wrapper = shallow(<Component />)
+
+  expect(wrapper.text()).toBe('empty')
+  wrapper.find('button').simulate('click')
+  expect(wrapper.text()).toBe('foobar')
+})
+
+test('withState (1 param) accepts setState() callback', () => {
+  const enhance = withState({ foo: 'foo' })
+  const callback = jest.fn()
+  const Component = enhance(({ setState }) =>
+    <button
+      onClick={() => setState({ foo: 'bar' }, callback)}
+    />
+  )
+
+  const wrapper = shallow(<Component />)
+
+  wrapper.find('button').simulate('click')
+  expect(callback).toHaveBeenCalled()
+})
+
+test('withState (1 param) accepts initialState as function of props', () => {
+  const enhance = withState(({ initialText }) => ({ text: initialText }))
+  const Component = enhance(({ state, setState }) =>
+    <button
+      onClick={() => setState({ text: 'new text' })}
+    >
+      {state.text}
+    </button>
+  )
+
+  const wrapper = shallow(<Component initialText="old text" />)
+
+  expect(wrapper.text()).toBe('old text')
+  wrapper.find('button').simulate('click')
+  expect(wrapper.text()).toBe('new text')
+})

src/packages/recompose/withState.js

@@ -2,7 +2,7 @@ import { Component } from 'react'
 import createHelper from './createHelper'
 import createEagerFactory from './createEagerFactory'
 
-const withState = (stateName, stateUpdaterName, initialState) =>
+const withStateOriginal = (stateName, stateUpdaterName, initialState) =>
   BaseComponent => {
     const factory = createEagerFactory(BaseComponent)
     return class extends Component {
@@ -30,4 +30,34 @@ const withState = (stateName, stateUpdaterName, initialState) =>
     }
   }
 
+const withClasslikeState = (initialState) =>
+  BaseComponent => {
+    const factory = createEagerFactory(BaseComponent)
+    return class extends Component {
+      state = typeof initialState === 'function'
+        ? initialState(this.props)
+        : initialState
+
+      setStateFn = (...args) => {
+        this.setState(...args)
+      }
+
+      render() {
+        return factory({
+          ...this.props,
+          state: this.state,
+          setState: this.setStateFn
+        })
+      }
+    }
+  }
+
+const withState = (...args) => {
+  if (args.length === 1) {
+    return withClasslikeState(...args)
+  }
+
+  return withStateOriginal(...args)
+}
+
 export default createHelper(withState, 'withState')