breaking out threat sources

rail/__init__.py

@@ -1,4 +1,6 @@
 from .cpi import CPI
 
 from .tree import Tree
+from .threat_source import ThreatSource, ThreatSources
 from .rail import ThreatSource, ThreatSources, ThreatEvent, ThreatEvents, Control, Controls, Vulnerability, Vulnerabilities, Likelihood, Impact, Risk, Risks
+

rail/rail.py

@@ -7,24 +7,13 @@
 from scipy.stats import lognorm
 
 from .tree import Tree
+from .threat_source import ThreatSource, ThreatSources
 
 pd.set_option('display.float_format', lambda x: '%.2f' % x)
 pd.set_option('display.max_colwidth', -1)
 plt.style.use('seaborn-poster')
 
 
-class ThreatSource(UserDict):
-    def __init__(self, name: str) -> None:
-        self.data = {}
-        self.data['name'] = name
-
-
-class ThreatSources(UserDict):
-    def new(self, name: str) -> ThreatSource:
-        self.data[name] = ThreatSource(name)
-        return self.data[name]
-
-
 class ThreatEvent(UserDict):
     def __init__(self, name: str, threat_source: ThreatSource) -> None:
         self.data = {}

rail/threat_source.py

@@ -0,0 +1,25 @@
+"""
+A class to implement Threat Sources
+"""
+from collections import UserDict
+
+class ThreatSource(UserDict): # pylint: disable=too-many-ancestors
+    """
+    A class to implement Threat Sources
+    """
+    def __init__(self, name: str) -> None:
+        UserDict.__init__(self)
+        self.data = {}
+        self.data['name'] = name
+
+
+class ThreatSources(UserDict): # pylint: disable=too-many-ancestors
+    """
+    A class to hold multiple Threat Sources
+    """
+    def new(self, name: str) -> ThreatSource:
+        """
+        A method to add a new threat source to the Threat Sources class
+        """
+        self.data[name] = ThreatSource(name)
+        return self.data[name]