# -*- coding: utf-8 -*-
"""common.py

Contains basic functions that are shared throughout the module
"""


def compute_totals(distribution, classes):
    "Compute the number of individuals per class, per unit and in total"
    N_unit = {au:sum([distribution[au][cl] for cl in classes]) for au in distribution}
    N_class = {cl:sum([dist_a[cl] for dist_a in distribution.values()]) for cl in classes}
    N_tot = sum(N_class.values())
    return N_unit, N_class, N_tot



def regroup_per_class(distribution, classes):
    "Return classes as they are presented in the data"

    try:
        new_distribution =  {au: {cl: sum([dist_au[c] for c in composition]) 
                                  for cl,composition in classes.iteritems()}
                             for au, dist_au in distribution.iteritems()}

    except KeyError:
        raise KeyError("Verify that the categories specified in the class"
                       " definitions exist in the original data.")

    return new_distribution



def return_categories(distribution):
    "Return the categories in the original data"
    keys = next(distribution.itervalues()).keys()
    return {k:[k] for k in keys}






# -*- coding: utf-8 -*-
"""common.py

Contains basic functions that are shared thoughout the module
"""


def compute_totals(distribution, classes):
    "Compute the number of individuals per class, per unit and in total"
    N_unit = {au:sum([distribution[au][cl] for cl in classes]) for au in distribution}
    N_class = {cl:sum([dist_a[cl] for dist_a in distribution.values()]) for cl in classes}
    N_tot = sum(N_class.values())
    return N_unit, N_class, N_tot



def regroup_per_class(distribution, classes):
    "Return classes as they are presented in the data"
    new_distribution =  {au: {cl: sum([dist_au[c] for c in composition]) 
                              for cl,composition in classes.iteritems()}
                         for au, dist_au in distribution.iteritems()}

    return new_distribution



def return_categories(distribution):
    "Return the categories in the original data"
    keys = next(distribution.itervalues()).keys()
    return {k:[k] for k in keys}






--- 
+++ 
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 """common.py
 
-Contains basic functions that are shared thoughout the module
+Contains basic functions that are shared throughout the module
 """
 
 
@@ -16,9 +16,15 @@
 
 def regroup_per_class(distribution, classes):
     "Return classes as they are presented in the data"
-    new_distribution =  {au: {cl: sum([dist_au[c] for c in composition]) 
-                              for cl,composition in classes.iteritems()}
-                         for au, dist_au in distribution.iteritems()}
+
+    try:
+        new_distribution =  {au: {cl: sum([dist_au[c] for c in composition]) 
+                                  for cl,composition in classes.iteritems()}
+                             for au, dist_au in distribution.iteritems()}
+
+    except KeyError:
+        raise KeyError("Verify that the categories specified in the class"
+                       " definitions exist in the original data.")
 
     return new_distribution
 
Raise exception if faulty definition of classes inserted
