From d733f429d98eae0e14e8ce3f3afc510ac43bcdbd Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Mon, 31 Oct 2022 19:09:59 +0800
Subject: [PATCH] feat(whl): check broken env before import

GitOrigin-RevId: 31ddb2a8c2504b8b98d1bf5ce0210cb61f637d59
---
 imperative/python/megengine/__init__.py | 32 ++++++++++++++++++++++++++++++++
 imperative/python/requires.txt          |  3 ++-
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/imperative/python/megengine/__init__.py b/imperative/python/megengine/__init__.py
index f71c96b9..8fd9e3bd 100644
--- a/imperative/python/megengine/__init__.py
+++ b/imperative/python/megengine/__init__.py
@@ -1,4 +1,36 @@
 # -*- coding: utf-8 -*-
+filter_package_name = 'megbrain'
+need_package_name = 'megengine'
+
+
+class MegEngineBorkenEnv(Exception):
+    pass
+
+
+import pkg_resources
+
+installed_packages = pkg_resources.working_set
+installed_packages_list = sorted(["%s" % i.key.lower() for i in installed_packages])
+is_broken_env = False
+if filter_package_name in installed_packages_list:
+    is_broken_env = True
+    print(
+        "broken env, please do uninstall {} and {}, then install {} again".format(
+            filter_package_name, need_package_name, need_package_name
+        )
+    )
+
+if "megenginelite" in installed_packages_list:
+    is_broken_env = True
+    print(
+        "broken env, please do uninstall megenginelite and {}, then install {} again".format(
+            need_package_name, need_package_name
+        )
+    )
+
+if is_broken_env:
+    raise MegEngineBorkenEnv
+
 import atexit
 import ctypes
 import re
diff --git a/imperative/python/requires.txt b/imperative/python/requires.txt
index d08566d3..e9e457d7 100644
--- a/imperative/python/requires.txt
+++ b/imperative/python/requires.txt
@@ -9,4 +9,5 @@ mprop
 wheel
 megfile>=0.0.10
 filelock
-psutil
\ No newline at end of file
+psutil
+setuptools