ban-module-import

// ban the module foo
// we can two types of imports
//  - import foo
//  - from foo import bar


// name of the packages to ban with their replacements
const PACKAGES_TO_BAN_WITH_REPLACEMENTS = new Map([
  ['foo', 'bar'],
  ['python2', null],
]);


// Report the violation and build the fix
const reportFooPackageUsage = (node, replacement) => {
  const error = buildError(node.start.line, node.start.col, node.end.line, node.end.col, "Cannot import foo", "CRITICAL", "SAFETY");
  
  if (replacement) {
    const edit = buildEditUpdate(node.start.line, node.start.col,
                                 node.end.line, node.end.col,
                                 replacement);
  	const fix = buildFix(`use ${replacement}`, [edit]);
  	addError(error.addFix(fix));
  } else {
    addError(error);
  }
};


// Main function
function visit(node, filename, code) {
  for(const name of PACKAGES_TO_BAN_WITH_REPLACEMENTS.keys()) {
    // catch the code "from <package> import <something>"
    if (node.astType === "fromstatement" && node.pkg) {
      if (node.pkg && node.pkg.value === name) {
        reportFooPackageUsage(node.pkg, PACKAGES_TO_BAN_WITH_REPLACEMENTS.get(name));
      }
  	}
    
    // catch the code "import <something>"
    if (node.astType === "importstatement" && node.packages) {
      const importFoo = node.packages.filter(p => p.name.value === name);
      if(importFoo.length > 0) {
        reportFooPackageUsage(importFoo[0], PACKAGES_TO_BAN_WITH_REPLACEMENTS.get(name));
      }
    }
  }

}