From a7f8f61fa26c727abe07b8733bed73f7ff118354 Mon Sep 17 00:00:00 2001
From: Semyon Proshev <Semyon.Proshev@jetbrains.com>
Date: Mon, 24 Jul 2017 16:49:12 +0300
Subject: [PATCH] Update setup.py creation to support importing setuptools
 instead of distutils.

---
 .../fileTemplates/internal/Setup Script.py.ft |  2 +-
 .../setupPy/CreateSetupPyAction.java          | 30 ++++++++++++++++---
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/python/resources/fileTemplates/internal/Setup Script.py.ft b/python/resources/fileTemplates/internal/Setup Script.py.ft
index cb9cb9b36e0b..c351cc60052d 100644
--- a/python/resources/fileTemplates/internal/Setup Script.py.ft	
+++ b/python/resources/fileTemplates/internal/Setup Script.py.ft	
@@ -1,4 +1,4 @@
-from distutils.core import setup
+$Import
 
 setup(
     name='$Package_name',
diff --git a/python/src/com/jetbrains/python/packaging/setupPy/CreateSetupPyAction.java b/python/src/com/jetbrains/python/packaging/setupPy/CreateSetupPyAction.java
index 749938ae36a8..86c87af505b6 100644
--- a/python/src/com/jetbrains/python/packaging/setupPy/CreateSetupPyAction.java
+++ b/python/src/com/jetbrains/python/packaging/setupPy/CreateSetupPyAction.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2014 JetBrains s.r.o.
+ * Copyright 2000-2017 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -27,6 +27,7 @@ import com.intellij.openapi.actionSystem.DataContext;
 import com.intellij.openapi.actionSystem.LangDataKeys;
 import com.intellij.openapi.module.Module;
 import com.intellij.openapi.project.Project;
+import com.intellij.openapi.projectRoots.Sdk;
 import com.intellij.openapi.roots.ModuleRootManager;
 import com.intellij.openapi.roots.ProjectFileIndex;
 import com.intellij.openapi.util.Comparing;
@@ -38,10 +39,16 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiManager;
 import com.intellij.util.SystemProperties;
 import com.jetbrains.python.PythonFileType;
+import com.jetbrains.python.packaging.PyPackage;
+import com.jetbrains.python.packaging.PyPackageManager;
 import com.jetbrains.python.packaging.PyPackageUtil;
 import com.jetbrains.python.psi.PyUtil;
+import com.jetbrains.python.sdk.PythonSdkType;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.util.Collection;
+import java.util.List;
 import java.util.Properties;
 
 /**
@@ -69,19 +76,34 @@ public class CreateSetupPyAction extends CreateFromTemplateAction {
 
   @Override
   public AttributesDefaults getAttributesDefaults(DataContext dataContext) {
-    Project project = CommonDataKeys.PROJECT.getData(dataContext);
+    final Project project = CommonDataKeys.PROJECT.getData(dataContext);
     final AttributesDefaults defaults = new AttributesDefaults("setup.py").withFixedName(true);
     if (project != null) {
+      defaults.addPredefined("Import", getSetupImport(dataContext));
       defaults.add("Package_name", project.getName());
       final PropertiesComponent properties = PropertiesComponent.getInstance();
-      defaults.add("Author", properties.getOrInit(AUTHOR_PROPERTY, SystemProperties.getUserName()));
-      defaults.add("Author_email", properties.getOrInit(EMAIL_PROPERTY, ""));
+      defaults.add("Author", properties.getValue(AUTHOR_PROPERTY, SystemProperties.getUserName()));
+      defaults.add("Author_email", properties.getValue(EMAIL_PROPERTY, ""));
       defaults.addPredefined("PackageList", getPackageList(dataContext));
       defaults.addPredefined("PackageDirs", getPackageDirs(dataContext));
     }
     return defaults;
   }
 
+  @NotNull
+  private static String getSetupImport(@NotNull DataContext dataContext) {
+    final Module module = LangDataKeys.MODULE.getData(dataContext);
+    return hasSetuptoolsPackage(module) ? "from setuptools import setup" : "from distutils.core import setup";
+  }
+
+  private static boolean hasSetuptoolsPackage(@Nullable Module module) {
+    final Sdk sdk = PythonSdkType.findPythonSdk(module);
+    if (sdk == null) return false;
+
+    final List<PyPackage> packages = PyPackageManager.getInstance(sdk).getPackages();
+    return packages != null && PyPackageUtil.findPackage(packages, "setuptools") != null;
+  }
+
   private static String getPackageList(DataContext dataContext) {
     final Module module = LangDataKeys.MODULE.getData(dataContext);
     if (module != null) {