From 9654455ee6677150cb198b261862cc5b04e90902 Mon Sep 17 00:00:00 2001
From: Vanya Belyaev <ivan.belyaev@cern.ch>
Date: Sat, 10 Aug 2024 13:41:11 +0200
Subject: [PATCH]   ?

---
 ostap/fitting/dataset.py                    | 40 ++++++++++++---------
 ostap/fitting/tests/test_fitting_dataset.py |  4 ++-
 2 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/ostap/fitting/dataset.py b/ostap/fitting/dataset.py
index db6fd2f6..98b2ff0c 100644
--- a/ostap/fitting/dataset.py
+++ b/ostap/fitting/dataset.py
@@ -86,7 +86,12 @@ def _rad_iter_ ( self ) :
 #  for index , entry, weight in dataset.loop ( 'pt>1' ) :
 #     print (index, entry, weight) 
 #  @endcode 
-def _rad_loop_ ( dataset , cuts = '' , cut_range = '' , first = 0 , last = LAST_ENTRY , progress = False ) :
+def _rad_loop_ ( dataset                ,
+                 cuts      = ''         ,
+                 cut_range = ''         ,
+                 first     = 0          ,
+                 last      = LAST_ENTRY ,
+                 progress = False       ) :
     """Iterator for `good' events  in dataset
     >>> dataset = ...
     >>> for index, entry, weight in dataset.loop ( ''pt>1' ) :
@@ -104,7 +109,7 @@ def _rad_loop_ ( dataset , cuts = '' , cut_range = '' , first = 0 , last = LAST_
     cut_range = str(cut_range).strip() if cut_range else ''
 
     fcuts = None 
-    if cuts : fcuts = make_formula ( cuts , cuts , dataset.varlist()   )
+    if cuts : fcuts = make_formula ( cuts , cuts , dataset.varlist() )
 
     weighted          = dataset.isWeighted                     ()
     store_errors      = weighted and dataset.store_errors      ()
@@ -118,16 +123,14 @@ def _rad_loop_ ( dataset , cuts = '' , cut_range = '' , first = 0 , last = LAST_
     nevents = 0 
     for event in source : 
 
-        vars , ww = dataset [ event ]        
-        if not vars : break
+        entry, weight = dataset [ event ]        
+        if not entry: break
 
-        if cut_range and not vars.allInRange ( cut_range ) : continue
+        if cut_range and not entry.allInRange ( cut_range ) : continue
 
         wc = fcuts.getVal() if fcuts else 1.0 
         if not wc  : continue
         
-        entry , weight = dataset [ event ]
-
         if weight is None : weight = wc if cuts else weight 
         else              : weight = weight * wc
 
@@ -203,26 +206,31 @@ def _rad_getitem_ ( data , index ) :
         if 1 == step and start <= stop : return data.reduce ( ROOT.RooFit.EventRange ( start , stop ) ) ## RETURN 
         index = range ( start , stop , step ) 
 
-    
-    ## "list" of indices
-    assert isinstance ( index , sequence_types ) , "Invalid type of `index':%s" % type ( index )
+    ## require sequence of indices hee 
+    if not isinstance ( index , sequence_types ) :
+        raise IndexError ( "Invalid type of `index':%s" % type ( index ) )
+
 
     weighted          = data.isWeighted                     ()
     store_errors      = weighted and data.store_errors      ()
     store_asym_errors = weighted and data.store_asym_errors ()
 
-    ## preare the resul 
+    ## preare the result 
     result = data.emptyClone ( dsID () )
     
-    ## the actual loop over entries 
+    ## the actual loop over set of entries 
     for i , j in enumerate ( index ) :
+
+        if not isinstance ( j , integer_types ) :
+            raise IndexError ( 'Invalid index [%s]=%s,' % ( i , j ) )
         
         j = int ( j )            ## the content must be convertible to integers 
-        if j < 0 : j += N        ## allow negative indices             
-        if not 0 <= j < N :      ## is adjusted integer in the proper range ?
+        if j < 0 : j += N        ## allow `slightly-negative' indices             
+        if not 0 <= j < N :      ## adjusted integer in the proper range ?
             raise IndexError ( 'Invalid index [%s]=%s,' % ( i , j ) ) 
-        
+
         vars = data.get ( j )
+        if not vars : raise IndexError ( 'Invalid index %s' % j )  ## 
         
         if   store_asym_errors :  
             wel , weh = data.weight_errors ()
@@ -240,7 +248,6 @@ def _rad_getitem_ ( data , index ) :
 # ==============================================================================
 
 
-
 # ==============================================================================
 ## Get (asymmetric) weigth errors for the current entry in dataset
 #  @code
@@ -1703,6 +1710,7 @@ def _rds_makeWeighted_ ( dataset       ,
                              cuts               ,
                              wvarname           )
 
+
 ROOT.RooDataSet.makeWeighted = _rds_makeWeighted_
 
 # =============================================================================
diff --git a/ostap/fitting/tests/test_fitting_dataset.py b/ostap/fitting/tests/test_fitting_dataset.py
index 37e83acc..623c6ef2 100644
--- a/ostap/fitting/tests/test_fitting_dataset.py
+++ b/ostap/fitting/tests/test_fitting_dataset.py
@@ -56,6 +56,9 @@
             
             mass.setVal ( random.uniform ( 0 , 10 ) )
             dataset.add ( varset )
+            
+
+""" 
 
 # =============================================================================
 weighted = dataset.makeWeighted ( 'Weight' )
@@ -66,7 +69,6 @@
 logger.info ( 'Print         unweighted dataset:\n%s' % dataset .table ( prefix = '# ' ) )
 logger.info ( 'Print           weighted dataset:\n%s' % weighted.table ( prefix = '# ' ) )
 
-""" 
 
 # =============================================================================
 ## (2) loop over some subset of entries