py_imports.grit

language python

/** This is a utility file, you do not need to implement it yourself */

pattern import_from($source, $names) {
    import_from_statement(name=$names, module_name=dotted_name(name=$source)),
}

pattern before_each_file_prep_imports() {
    $_ where {
        $GLOBAL_NEW_BARE_IMPORTS = [],
        $GLOBAL_NEW_FROM_IMPORTS = [],
        $GLOBAL_IMPORTS_TO_REMOVE = [],
    }
}

pattern after_each_file_handle_imports() {
  file($body) where {
    $body <: maybe insert_imports(),
    $body <: maybe handle_remove_imports()
  }
}

pattern process_one_source($all_imports) {
    [$program, $source, $new_names] where {
        if ($program <: module(statements = some import_from($source, $names))) {
            $names_to_add = "",
            $new_names <: some $name where {
                if (!$names <: some $name) {
                    $names_to_add += `, $name`
                }
            },
            $names => `$names$names_to_add`,
        } else {
            $joined_names = join(list = $new_names, separator = ", "),
            $all_imports += `from $source import $joined_names\n`,
        }
    }
}

// TODO: remove imports from the global list as we remove them
pattern handle_one_removal_candidate() {
  // $name represents the AST node we imported from a package that should be removed
  or {
    `from $src import $mods` as $full_import where {
      $keep_some = false,
      $mods <: some bubble($keep_some) or {
        aliased_import(alias=$name) where {
          $GLOBAL_IMPORTS_TO_REMOVE <: some bubble($name) $name
        } => .,
        `$name` where {
          $GLOBAL_IMPORTS_TO_REMOVE <: some bubble($name) $name
        } => .,
        $other where {
          $keep_some = true
        }
      },
      if ($keep_some <: false) {
        $full_import => .
      }
    },
    `import $mod` where {
      $GLOBAL_IMPORTS_TO_REMOVE <: some bubble($mod) $mod
    } => .,
    `import $mod as $alias` where {
      $GLOBAL_IMPORTS_TO_REMOVE <: some bubble($alias) $alias
    } => .,
  }
}

pattern handle_remove_imports() {
  $body where {
    $body <: contains handle_one_removal_candidate()
  }
}

pattern remove_import() {
  $name where {
    if ($GLOBAL_IMPORTS_TO_REMOVE <: not some $name) {
      $GLOBAL_IMPORTS_TO_REMOVE += $name
    }
  }
}

pattern insert_imports() {
    $body where {
        $all_imports = "",
        $GLOBAL_NEW_FROM_IMPORTS <: maybe some process_one_source($all_imports),
        $GLOBAL_NEW_BARE_IMPORTS <: maybe some bubble($all_imports) $name where {
            $all_imports += `import $name\n`,
        },
        if (!$all_imports <: "") {
            $body => `$all_imports\n$body`
        }
    }
}

pattern imported_from($source) {
    $name where {
        $program <: module($statements),
        $statements <: some bubble($name, $source) import_from($source, $names) where {
            $names <: some $name,
        }
    }
}

pattern ensure_import_from($source) {
    $name where {
        if ($name <: not imported_from($source)) {
            or {
              $GLOBAL_NEW_FROM_IMPORTS <: some [$program, $source, $existing] where {
                $existing += $name
              },
              $GLOBAL_NEW_FROM_IMPORTS += [$program, $source, [$name]]
            }
        }
    }
}

pattern is_bare_imported() {
    $name where {
        $program <: module($statements),
        $statements <: some import_statement(name=$names) where {
            $names <: some dotted_name(name=$name),
        },
    }
}

pattern ensure_bare_import() {
    $name where {
        if ($name <: not is_bare_imported()) {
            if ($GLOBAL_NEW_BARE_IMPORTS <: not some $name) {
                $GLOBAL_NEW_BARE_IMPORTS += [$name]
            }
        }
    }
}

predicate add_import($name, $source) {
    if ($name <: undefined) {
      $source <: ensure_bare_import(),
    } else {
      $name <: ensure_import_from(source=$source)
    }
}