From 7f39c7721995c66811b7ebb90a86fbba740140ef Mon Sep 17 00:00:00 2001
From: ggrieco-tob <gustavo.grieco@trailofbits.com>
Date: Fri, 10 May 2019 12:23:51 -0300
Subject: [PATCH] improved module loading

---
 utils/similarity/__init__.py | 10 +---------
 utils/similarity/info.py     |  2 +-
 utils/similarity/model.py    | 11 +++++++++++
 utils/similarity/plot.py     |  6 +++---
 utils/similarity/test.py     |  6 +++---
 utils/similarity/train.py    |  2 +-
 6 files changed, 20 insertions(+), 17 deletions(-)
 create mode 100644 utils/similarity/model.py

diff --git a/utils/similarity/__init__.py b/utils/similarity/__init__.py
index e8ce679e9..b31b92c60 100644
--- a/utils/similarity/__init__.py
+++ b/utils/similarity/__init__.py
@@ -1,9 +1 @@
-# from https://stackoverflow.com/questions/563022/whats-python-good-practice-for-importing-and-offering-optional-features
-import sys
-
-try:
-    import fastText
-except ImportError:
-    print("ERROR: in order to use slither-simil, you need to install fastText 0.2.0:")
-    print("$ pip3 install https://github.com/facebookresearch/fastText/archive/0.2.0.zip --user\n")
-    sys.exit(-1)
+from .model import load_model
diff --git a/utils/similarity/info.py b/utils/similarity/info.py
index 431852571..e250aa991 100644
--- a/utils/similarity/info.py
+++ b/utils/similarity/info.py
@@ -3,7 +3,7 @@ import sys
 import os.path
 import traceback
 
-from fastText import load_model
+from .model  import load_model
 from .encode import parse_target, encode_contract
 
 logging.basicConfig()
diff --git a/utils/similarity/model.py b/utils/similarity/model.py
new file mode 100644
index 000000000..5e1549058
--- /dev/null
+++ b/utils/similarity/model.py
@@ -0,0 +1,11 @@
+import sys
+
+try:
+    from fastText import load_model
+    from fastText import train_unsupervised
+except ImportError:
+    print("ERROR: in order to use slither-simil, you need to install fastText 0.2.0:")
+    print("$ pip3 install https://github.com/facebookresearch/fastText/archive/0.2.0.zip --user\n")
+    sys.exit(-1)
+
+
diff --git a/utils/similarity/plot.py b/utils/similarity/plot.py
index 1b4d07e02..05d8bf921 100644
--- a/utils/similarity/plot.py
+++ b/utils/similarity/plot.py
@@ -5,15 +5,15 @@ import operator
 import numpy as np
 import random
 
+from .model  import load_model   
+from .encode import load_and_encode, parse_target
+
 try:
     from sklearn import decomposition
     import matplotlib.pyplot as plt
 except ImportError:
     decomposition = None
     plt = None
-    
-from fastText import load_model
-from .encode import load_and_encode, parse_target
 
 logger = logging.getLogger("Slither-simil")
 
diff --git a/utils/similarity/test.py b/utils/similarity/test.py
index f821b6448..08542dd0d 100755
--- a/utils/similarity/test.py
+++ b/utils/similarity/test.py
@@ -5,9 +5,9 @@ import traceback
 import operator
 import numpy as np
 
-from fastText import load_model
-from .encode import encode_contract, load_and_encode, parse_target
-from .cache import save_cache
+from .model      import load_model
+from .encode     import encode_contract, load_and_encode, parse_target
+from .cache      import save_cache
 from .similarity import similarity
 
 logger = logging.getLogger("Slither-simil")
diff --git a/utils/similarity/train.py b/utils/similarity/train.py
index d3039be38..e810450a6 100755
--- a/utils/similarity/train.py
+++ b/utils/similarity/train.py
@@ -5,7 +5,7 @@ import traceback
 import operator
 import os
 
-from fastText import train_unsupervised
+from .model   import train_unsupervised
 from .encode  import encode_contract, load_contracts
 from .cache   import save_cache