Changeset 1166

Show
Ignore:
Timestamp:
07/28/08 18:45:08 (4 months ago)
Author:
pierregm
Message:

* tseries : rationalize .getitem

Files:

Legend:

Unmodified
Added
Removed
Modified
Copied
Moved

  • trunk/timeseries/scikits/timeseries/tests/test_timeseries.py

    r1165 r1166  
    505505 
    506506 
     507    def test_slicing_and_keeping_additional_attributes(self): 
     508        series1D = self.series1D 
     509        series1D.fill_value = -9999 
     510        series1D._basedict['info'] = '???' 
     511        piece = series1D[:5] 
     512        assert_equal(piece._fill_value, -9999) 
     513        assert_equal(piece[:5]._basedict['info'], '???') 
     514         
     515 
    507516class TestSetItem(TestCase): 
    508517    # 

  • trunk/timeseries/scikits/timeseries/tseries.py

    r1165 r1166  
    506506        """ 
    507507        (sindx, dindx, recheck) = self._index_checker(indx) 
    508         _class = self.__class__ 
     508        _data = ndarray.__getattribute__(self, '_data') 
     509        _mask = ndarray.__getattribute__(self, '_mask') 
    509510        _dates = ndarray.__getattribute__(self, '_dates') 
    510511        try: 
    511             # temporarily change class to MaskedArray for indexing. MUCH 
    512             # faster than using a view due to complexity of __array_finalize__ 
    513             # for MaskedArray class 
    514             self.__class__ = MaskedArray 
    515             newseries = self.__getitem__(sindx) 
     512            output = _data.__getitem__(sindx) 
    516513        except IndexError: 
    517514            # We don't need to recheck the index: just raise an exception 
    518515            if not recheck: 
    519                 self.__class__ = _class 
    520516                raise 
    521517            # Maybe the index is a list of Dates ? 
    …… 
    530526                    exc_info = sys.exc_info() 
    531527                    msg = "Invalid index or date '%s'" % indx 
    532                     self.__class__ = _class 
    533528                    raise IndexError(msg), None, exc_info[2] 
    534529                else: 
    535                     newseries = self.__getitem__(indx) 
    536                     dindx = indx 
    537                 self.__class__ = _class 
     530                    output = _data.__getitem__(indx) 
     531                    sindx = dindx = indx 
    538532            else: 
    539                 newseries = self.__getitem__(indx) 
    540                 dindx = indx 
    541                 self.__class__ = _class 
    542         self.__class__ = _class 
     533                output = _data.__getitem__(indx) 
     534                sindx = dindx = indx 
    543535        # Don't find the date if it's not needed...... 
    544         if np.isscalar(newseries) or (newseries is masked): 
    545             return newseries 
     536        if not getattr(output,'ndim', False): 
     537            if _mask is not nomask and _mask[sindx]: 
     538                return masked 
     539            return output 
    546540        # Get the date................................ 
    547541        newdates = _dates.__getitem__(dindx) 
    548         # In fact, that's a scalar 
    549542        if not getattr(newdates, 'shape', 0): 
    550             return newseries 
    551         newseries = newseries.view(type(self)) 
    552         newseries._dates = newdates 
     543            # No dates ? Output a MaskedArray 
     544            newseries = output.view(MaskedArray) 
     545        else: 
     546            # Some dates: output a TimeSeries 
     547            newseries = output.view(type(self)) 
     548            newseries._dates = newdates 
     549        # Update some info from self (fill_value, _basedict...) 
     550        MaskedArray._update_from(newseries, self) 
     551        # Fix the fill_value if we were accessing a field of a flexible array 
     552        if isinstance(indx, basestring): 
     553            if self._fill_value is not None: 
     554                 newseries._fill_value = self._fill_value[indx] 
     555            newseries._isfield = True 
     556        # Fix the mask 
     557        if _mask is not nomask: 
     558            newseries._mask = _mask[sindx] 
     559            newseries._sharedmask = True 
    553560        return newseries 
    554561