contextTypeHandler and childContextTypeHandler

src/Documentation.js

@@ -12,11 +12,15 @@
 
 class Documentation {
   _props: Object;
+  _context: Object;
+  _childContext: Object;
   _composes: Set<string>;
   _data: Object;
 
   constructor() {
     this._props = new Map();
+    this._context = new Map();
+    this._childContext = new Map();
     this._composes = new Set();
     this._data = new Map();
   }
@@ -41,6 +45,22 @@ class Documentation {
     return propDescriptor;
   }
 
+  getContextDescriptor(propName: string): PropDescriptor {
+    var propDescriptor = this._context.get(propName);
+    if (!propDescriptor) {
+      this._context.set(propName, propDescriptor = {});
+    }
+    return propDescriptor;
+  }
+
+  getChildContextDescriptor(propName: string): PropDescriptor {
+    var propDescriptor = this._childContext.get(propName);
+    if (!propDescriptor) {
+      this._childContext.set(propName, propDescriptor = {});
+    }
+    return propDescriptor;
+  }
+
   toObject(): Object {
     var obj = {};
 
@@ -55,6 +75,20 @@ class Documentation {
       }
     }
 
+    if (this._context.size > 0) {
+      obj.context = {};
+      for (var [name, descriptor] of this._context) {
+        obj.context[name] = descriptor;
+      }
+    }
+
+    if (this._childContext.size > 0) {
+      obj.childContext = {};
+      for (var [name, descriptor] of this._childContext) {
+        obj.childContext[name] = descriptor;
+      }
+    }
+
     if (this._composes.size > 0) {
       obj.composes = Array.from(this._composes);
     }

src/handlers/__tests__/propTypeHandler-test.js

@@ -15,7 +15,7 @@ jest.mock('../../utils/getPropType', () => jest.fn(() => ({})));
 
 import {statement, expression} from '../../../tests/utils';
 import Documentation from '../../Documentation';
-import propTypeHandler from '../propTypeHandler';
+import {propTypeHandler} from '../propTypeHandler';
 
 describe('propTypeHandler', () => {
   var getPropTypeMock;

src/handlers/index.js

@@ -15,7 +15,7 @@ export {default as componentDocblockHandler} from './componentDocblockHandler';
 export {default as componentMethodsHandler} from './componentMethodsHandler';
 export {default as componentMethodsJsDocHandler} from './componentMethodsJsDocHandler';
 export {default as defaultPropsHandler} from './defaultPropsHandler';
-export {default as propTypeHandler} from './propTypeHandler';
+export {propTypeHandler, contextTypeHandler, childContextTypeHandler} from './propTypeHandler';
 export {default as propTypeCompositionHandler} from './propTypeCompositionHandler';
 export {default as propDocBlockHandler} from './propDocBlockHandler';
 export {default as displayNameHandler} from './displayNameHandler';

src/handlers/propTypeHandler.js

@@ -33,15 +33,15 @@ function isPropTypesExpression(path) {
   return false;
 }
 
-function amendPropTypes(documentation, path) {
+function amendPropTypes(getDescriptor, path) {
   if (!types.ObjectExpression.check(path.node)) {
     return;
   }
 
   path.get('properties').each(function(propertyPath) {
     switch (propertyPath.node.type) {
       case types.Property.name:
-        var propDescriptor = documentation.getPropDescriptor(
+        var propDescriptor = getDescriptor(
           getPropertyName(propertyPath)
         );
         var valuePath = propertyPath.get('value');
@@ -59,25 +59,42 @@ function amendPropTypes(documentation, path) {
         var resolvedValuePath = resolveToValue(propertyPath.get('argument'));
         switch (resolvedValuePath.node.type) {
           case types.ObjectExpression.name: // normal object literal
-            amendPropTypes(documentation, resolvedValuePath);
+            amendPropTypes(getDescriptor, resolvedValuePath);
             break;
         }
         break;
     }
   });
 }
 
-export default function propTypeHandler(
-  documentation: Documentation,
-  path: NodePath
-) {
-  var propTypesPath = getMemberValuePath(path, 'propTypes');
-  if (!propTypesPath) {
-    return;
-  }
-  propTypesPath = resolveToValue(propTypesPath);
-  if (!propTypesPath) {
-    return;
+export function getPropTypeHandler(propName: string) {
+  return function (
+    documentation: Documentation,
+    path: NodePath
+  ) {
+    var propTypesPath = getMemberValuePath(path, propName);
+    if (!propTypesPath) {
+      return;
+    }
+    propTypesPath = resolveToValue(propTypesPath);
+    if (!propTypesPath) {
+      return;
+    }
+    let getDescriptor;
+    switch(propName) {
+      case 'childContextTypes':
+        getDescriptor = documentation.getChildContextDescriptor;
+        break;
+      case 'contextTypes':
+        getDescriptor = documentation.getContextDescriptor;
+        break;
+      default:
+        getDescriptor = documentation.getPropDescriptor;
+    }
+    amendPropTypes(getDescriptor.bind(documentation), propTypesPath);
   }
-  amendPropTypes(documentation, propTypesPath);
 }
+
+export const propTypeHandler = getPropTypeHandler('propTypes')
+export const contextTypeHandler = getPropTypeHandler('contextTypes')
+export const childContextTypeHandler = getPropTypeHandler('childContextTypes')