babel-plugin-transform-private-to-weakmap/src/index.js

Transforms class properties prefixed with an underscore to weakmaps.

export default function({ types: t }) {
  return {
    visitor: {
      Program(programPath) {
        let propsToRemove = [];
        programPath.traverse({
          // Find all classes and test them for private properties
          Class(classPath) {
            let classPrivateProps = {};

            let insertPath =
              classPath.parentPath.type === 'ExportNamedDeclaration' ||
              classPath.parentPath.type === 'ExportDefaultDeclaration'
              ? classPath.parentPath
              : classPath
            ;
            let constructor = null;

            let hasPrivate = false;

            // let body = classPath.get('body');
            // for(let subPath of body.get('body')) console.log(subPath.type);

            classPath.traverse({
              ClassProperty(propPath) {
                if(propPath.node.key.name.startsWith('_')) {
                  hasPrivate = true;
                  classPrivateProps[propPath.node.key.name] = propPath.node.value;
                  propsToRemove.push(propPath);
                }
              }
            });

            if(hasPrivate) {
              // Transform private access into weakmap access
              classPath.traverse({
                AssignmentExpression(assignPath) {
                  handleAssignPrivateProperties(t, assignPath, classPrivateProps, classPath);
                },
                Expression(exprPath) {
                  if(exprPath.type === 'MemberExpression')
                    handleAccessPrivateProperties(t, exprPath, classPrivateProps, classPath);
                },

                ClassMethod(methodPath) {
                  if(methodPath.node.kind === 'constructor') constructor = methodPath;

                  methodPath.traverse({
                    AssignmentExpression(assignPath) {
                      handleAssignPrivateProperties(t, assignPath, classPrivateProps, classPath);
                    },
                    Expression(exprPath) {
                      if(exprPath.type === 'MemberExpression')
                        handleAccessPrivateProperties(t, exprPath, classPrivateProps, classPath);
                    }
                  });
                }
              });

              // Insert the weakmap definition in the class
              insertPath.insertAfter(
                t.expressionStatement(
                  t.assignmentExpression(
                    '=',
                    t.memberExpression(
                      classPath.node.id,
                      t.identifier('private'),
                      false,
                      false
                    ),
                    t.newExpression(t.identifier('WeakMap'), [])
                  )
                )
              );

              // Insert an empty object initialization for the weakmap
              // in the constructor

              // Generate object with private props that have a default value
              let i, objProps=[], objPropNames = Object.keys(classPrivateProps);
              for(i=0; i<objPropNames.length; i++) {
                if(classPrivateProps[objPropNames[i]] !== null)
                  objProps.push(t.objectProperty(
                    t.identifier(objPropNames[i].substr(1)),
                    classPrivateProps[objPropNames[i]]
                  ));
              }
              let obj = t.objectExpression(objProps);

              // Generate the weakmap .set() call with the generated object
              let privateInit = t.expressionStatement(
                t.callExpression(
                  t.memberExpression(
                    t.memberExpression(
                      classPath.node.id,
                      t.identifier('private'),
                      false,
                      false
                    ),
                    t.identifier('set'),
                    false,
                    false
                  ),
                  [ t.thisExpression(), obj ]
                )
              );

              // Append to existing constructor
              if(constructor !== null) {
                constructor.get('body').unshiftContainer('body', privateInit );
              }

              // or create a constructor if there was none
              else {
                classPath.get('body').unshiftContainer('body', t.classMethod(
                  'constructor',
                  t.identifier('constructor'),
                  [],
                  t.blockStatement([ privateInit ])
                ))
              }
            }
          }
        });
        propsToRemove.forEach(prop => prop.remove());
      }
    }
  };
}

const getPrivatePropName = (path, classPrivateProps) => {
  let prop = null;
  if(path.node.type === 'MemberExpression') {
    if(path.node.object.type === 'ThisExpression') {
      if(path.node.property.type === 'Identifier') {
        if ('_' + path.node.property.name in classPrivateProps)
          prop = path.node.property.name;
      }
    }
  }
  return prop;
};

const getAssignPrivatePropName = (path, classPrivateProps) => {
  let prop = null;
  if(path.node.left.object.type !== 'ThisExpression') return;
  if(path.node.left.type === 'MemberExpression') {
    if(path.node.left.property.type === 'Identifier') {
      if ('_' + path.node.left.property.name in classPrivateProps)
        prop = path.node.left.property.name;
    }
  }
  else if(path.node.left.name in classPrivateProps) {
    prop = path.node.left.name;
  }

  return prop;
};

const handleAccessPrivateProperties = (t, path, classPrivateProps, classPath) => {
  let prop = getPrivatePropName(path, classPrivateProps);
  if(prop !== null && typeof prop !== 'undefined') {
    path.replaceWith(t.expressionStatement(
      t.memberExpression(
        t.callExpression(
          t.memberExpression(
            t.memberExpression(
              classPath.node.id,
              t.identifier('private'),
              false,
              false
            ),
            t.identifier('get'),
            false,
            false
          ),
          [t.thisExpression()]
        ),
        t.identifier(prop),
        false,
        false
      )
    ));
  }
};

const handleAssignPrivateProperties = (t, path, classPrivateProps, classPath) => {
  let prop = getAssignPrivatePropName(path, classPrivateProps);
  if(prop !== null && typeof prop !== 'undefined') {
    path.replaceWith(t.expressionStatement(
      t.assignmentExpression(
        path.node.operator,
        t.memberExpression(
          t.callExpression(
            t.memberExpression(
              t.memberExpression(
                classPath.node.id,
                t.identifier('private'),
                false,
                false
              ),
              t.identifier('get'),
              false,
              false
            ),
            [t.thisExpression()]
          ),
          t.identifier(prop),
          false,
          false
        ),
        path.node.right
      )
    ));
  }
};