withStyle-babel-plugin.js

// const { default: generate } = require('@babel/generator')
module.exports = function ({ types: t }) {
  let programPath
  function getUniqueName (key = '__style') {
    return programPath.scope.generateUidIdentifier(key).name
  }

  function getStyles (body, _this) {
    let namedImports = []
    let unnamedImports = []
    let numImports = 0
    let withStyleImported = false
    for (let i = 0; i < body.length; i++) {
      const item = body[i]
      if (item.type === 'ImportDeclaration') {
        numImports++
        if (item.source.value === 'style-loader') withStyleImported = true

        if (item.source.value.includes('.css')) {
          let name
          for (let specifier of item.specifiers) {
            if (specifier.type === 'ImportDefaultSpecifier') {
              name = specifier.local.name
              break
            }
          }
          if (!name) {
            name = getUniqueName()
            item.specifiers.unshift(
              t.importDefaultSpecifier(t.identifier(name))
            )
            unnamedImports.push(t.identifier(name))
          } else {
            _this.namedImports.push(name)
            namedImports.push(t.identifier(name))
          }
        }
      }

      if (
        item.type === 'ExpressionStatement' &&
        item.expression &&
        item.expression.type === 'CallExpression' &&
        item.expression.callee &&
        item.expression.callee.name === 'require' &&
        item.expression.arguments[0].value.includes('.css')
      ) {
        let name
        name = getUniqueName()
        body.splice(i, 1)
        body.splice(
          numImports,
          0,
          t.importDeclaration(
            [t.importDefaultSpecifier(t.identifier(name))],
            t.stringLiteral(item.expression.arguments[0].value)
          )
        )
        numImports++
        unnamedImports.push(t.identifier(name))
      }
    }
    let addUnNamedImports = false
    let addNamedImports = false
    let styleNames = []
    let imports = []
    if (namedImports.length) {
      addNamedImports = true
      imports.push(t.identifier('__namedStyleImports'))
      styleNames.push(t.spreadElement(t.identifier('__namedStyleImports')))
      body.splice(
        numImports,
        0,
        t.variableDeclaration('const', [
          t.variableDeclarator(
            t.identifier('__namedStyleImports'),
            t.arrayExpression(namedImports)
          )
        ])
      )
      numImports++
    }
    if (unnamedImports.length) {
      addUnNamedImports = true
      imports.push(t.identifier('__unnamedStyleImports'))
      styleNames.push(t.spreadElement(t.identifier('__unnamedStyleImports')))
      body.splice(
        numImports,
        0,
        t.variableDeclaration('const', [
          t.variableDeclarator(
            t.identifier('__unnamedStyleImports'),
            t.arrayExpression(unnamedImports)
          )
        ])
      )
      numImports++
    }
    ;(addUnNamedImports || addNamedImports) &&
      body.splice(
        numImports,
        0,
        t.exportNamedDeclaration(
          t.variableDeclaration('const', [
            t.variableDeclarator(
              t.identifier('__styles'),
              t.arrayExpression(styleNames)
            )
          ])
        )
      )

    _this.addUnNamedImports = addUnNamedImports
    _this.addNamedImports = addNamedImports
    return { withStyleImported, value: imports.length ? imports : undefined }
  }

  function processImportedStyles (programPath, body, localNames, map) {
    let styleRegex = /(style|Style|style.jsx|Style.jsx)$/
    
    let names = []
    if (localNames) {
      names = names.concat(localNames)
    }
    const itemsToAdd = []
    for (let item of body) {
      if (
        item.type === 'ImportDeclaration' &&
        styleRegex.test(item.source.value) &&
        !item.source.value.includes('.css')
      ) {
        let name = getUniqueName('_importedStyle')
        for (let specifier of item.specifiers) {
          const { local: { name: localName } = {} } = specifier
          if (localName) {
            map[localName] = name
          }
        }
        itemsToAdd.push(
          t.importDeclaration(
            [t.importSpecifier(t.identifier(name), t.identifier('__styles'))],
            item.source
          )
        )
        
        names.push(t.identifier(name))
      }
    }
    programPath.node.body = itemsToAdd.concat(body)
    if (names.length) {
      return t.arrayExpression(names)
    }
  }

  function wrapDefaultExport (body, styleArray, withStyleImported) {
    if (styleArray && !withStyleImported) {
      const exportDefault = body.find(
        item => item.type === 'ExportDefaultDeclaration'
      )
      const { type, callee: { name } = {} } =
        (exportDefault || {}).declaration || {}
      if (
        exportDefault &&
        (type !== 'CallExpression' || name !== 'withStyles')
      ) {
        body.unshift(
          t.importDeclaration(
            [t.importDefaultSpecifier(t.identifier('withStyles'))],
            t.stringLiteral('style-loader')
          )
        )
        if (exportDefault.declaration.type === 'FunctionDeclaration') {
          exportDefault.declaration.type = 'FunctionExpression'
        }
        exportDefault.declaration = t.callExpression(
          t.identifier('withStyles'),
          [exportDefault.declaration, styleArray]
        )
      }
    }
  }

  return {
    pre () {
      this.map = {}
      this.namedImports = []
    },
    visitor: {
      Program (path, state) {
        // ;(state.file.opts.filename.includes('test.jsx') ||
        //   state.file.opts.filename.includes('testStyle.jsx')) &&
        //   console.log(generate(path.parent).code)
        programPath = path
        const map = this.map
        const { value: names, withStyleImported } = getStyles(
          path.node.body,
          this
        )
        const processedStyles = processImportedStyles(
          path,
          path.node.body,
          names,
          map,
          state
        )
        wrapDefaultExport(path.node.body, processedStyles, withStyleImported)
      },
      CallExpression (path) {
        if (
          path.node.callee &&
          ['withStyles', 'useStyles'].includes(path.node.callee.name)
        ) {
          const index = path.node.callee.name === 'useStyles' ? 0 : 1
          const array = path.node.arguments[index].elements || []
          path.node.arguments[index].elements = array
          let unNamedAdded = false
          for (let i = 0; i < array.length; i++) {
            const { name } = array[i]
            if (name === '__unnamedStyleImports') {
              unNamedAdded = true
            }
            if (this.namedImports.includes(name)) {
              array[i] = t.arrayExpression([array[i]])
            } else {
              array[i] = t.identifier(this.map[name] || name)
            }
          }
          !unNamedAdded &&
            this.addUnNamedImports &&
            array.push(t.identifier('__unnamedStyleImports'))
        }
      }
    }
  }
}