Merge branch 'enable-pruning-in-crosstabs-#158458376' into rel-5.2.102

src/cr/cube/crunch_cube.py

@@ -1083,34 +1083,6 @@ def population_counts(self, population_size, weighted=True,
             prune=prune
         ) * population_size
 
-    def y_offset(self, expand=False, include_transforms_for_dims=None):
-        '''Gets y offset for sheet manipulation.
-
-        Args:
-            - expand (bool): If a cube is a categorical array, and it's also a
-                0-index cube of multiple cube sequence, it needs to be
-                displayed in a specific manner (sliced across categories
-                dimension). This argument enables such behavior. If a CA cube
-                is the only one (not a part of the multitable), it is output
-                normally, without expansion.
-        '''
-        if not self.dimensions:
-            return 4
-
-        first_dim_length = self.as_array(
-            include_transforms_for_dims=include_transforms_for_dims
-        ).shape[0]
-
-        # Special case of CA as a 0-ind cube
-        if expand and self.dimensions[0].type == 'categorical_array':
-            return first_dim_length * (len(self.dimensions[1].elements()) + 4)
-
-        if self.ndim <= 2 and self.dimensions[0].type:
-            return first_dim_length + 4
-        return first_dim_length * (self.as_array(
-            include_transforms_for_dims=include_transforms_for_dims
-        ).shape[1] + 4)
-
     def index(self, weighted=True, prune=False):
         '''Get cube index measurement.'''
         return Index(self, weighted, prune).data

src/cr/cube/cube_slice.py

@@ -150,3 +150,33 @@ def ca_main_axis(self):
             return 1 - ca_ind
         else:
             return None
+
+    def labels(self, hs_dims=None, prune=False):
+        '''Get labels for the cube slice, and perform pruning by slice.'''
+        labels = self._cube.labels(include_transforms_for_dims=hs_dims)[-2:]
+        if not prune:
+            return labels
+
+        def prune_dimension_labels(labels, prune_indices):
+            '''Get pruned labels for single dimension, besed on prune inds.'''
+            labels = [
+                label for label, prune in zip(labels, prune_indices)
+                if not prune
+            ]
+            return labels
+
+        labels = [
+            prune_dimension_labels(dim_labels, dim_prune_inds)
+            for dim_labels, dim_prune_inds in
+            zip(labels, self.prune_indices(transforms=hs_dims))
+        ]
+        return labels
+
+    @property
+    def has_mr(self):
+        '''True if the slice has MR dimension.
+
+        This property needs to be overridden, because we don't care about the
+        0th dimension (and if it's an MR) in the case of a 3D cube.
+        '''
+        return 'multiple_response' in self.dim_types

src/cr/cube/subtotal.py

@@ -20,8 +20,11 @@ def is_valid(self):
         if isinstance(self._data, dict):
             required_keys = {'anchor', 'args', 'function', 'name'}
             has_keys = set(self._data.keys()) == required_keys
-            if has_keys:
-                return self._data['function'] == 'subtotal'
+            if has_keys and self._data['function'] == 'subtotal':
+                return any(
+                    element for element in self._dim.elements()
+                    if element['id'] in self._data['args']
+                )
         return False
 
     @lazyproperty
@@ -46,8 +49,9 @@ def _all_dim_ids(self):
     @lazyproperty
     def args(self):
         '''Get H&S args.'''
-        if self.is_valid:
-            return self._data['args']
+        hs_ids = self._data.get('args', None)
+        if hs_ids and self.is_valid:
+            return hs_ids
         return []
 
     @lazyproperty

tests/unit/test_crunch_cube.py

@@ -194,59 +194,6 @@ def test_test_filter_annotation(self):
         actual = CrunchCube(mock_cube).filter_annotation
         self.assertEqual(actual, expected)
 
-    @patch('cr.cube.crunch_cube.CrunchCube.dimensions')
-    def test_y_offset_no_dimensions(self, mock_dims):
-        dims = []
-        mock_dims.__get__ = Mock(return_value=dims)
-        expected = 4
-        actual = CrunchCube({}).y_offset()
-        self.assertEqual(actual, expected)
-
-    @patch('cr.cube.crunch_cube.CrunchCube.as_array')
-    @patch('cr.cube.crunch_cube.CrunchCube.dimensions')
-    def test_y_offset_two_dimensions(self, mock_dims, mock_as_array):
-        dims = [Mock(), Mock()]
-        mock_dims.__get__ = Mock(return_value=dims)
-        mock_as_array.return_value = np.arange(12).reshape(3, 4)
-
-        # Expected is the first dim length + offset of 4
-        expected = 3 + 4
-        actual = CrunchCube({}).y_offset()
-        self.assertEqual(actual, expected)
-
-    @patch('cr.cube.crunch_cube.CrunchCube.as_array')
-    @patch('cr.cube.crunch_cube.CrunchCube.dimensions')
-    def test_y_offset_three_dimensions(self, mock_dims, mock_as_array):
-        dims = [Mock(), Mock(), Mock()]
-        mock_dims.__get__ = Mock(return_value=dims)
-        mock_as_array.return_value = np.arange(24).reshape(3, 4, 2)
-
-        # Expected is the first dim len * (sedond dim len + offset of 4)
-        # The first dim is used as a slice, while the second is the number
-        # of rows of each subtable. Offset is space between tables
-        expected = 3 * (4 + 4)
-        actual = CrunchCube({}).y_offset()
-        self.assertEqual(actual, expected)
-
-    @patch('cr.cube.crunch_cube.CrunchCube.as_array')
-    @patch('cr.cube.crunch_cube.CrunchCube.dimensions')
-    def test_y_offset_ca_as_zero_index_cube(self, mock_dims, mock_as_array):
-        dims = [Mock(), Mock()]
-        dims[0].type = 'categorical_array'
-
-        # Second dim has total length of 3
-        dims[1].elements.return_value = [Mock(), Mock(), Mock()]
-        mock_dims.__get__ = Mock(return_value=dims)
-        mock_as_array.return_value = np.arange(12).reshape(3, 4)
-
-        # Expected is the first dim len * (sedond dim len + offset of 4)
-        # The first dim is used as a slice, while the second is the number
-        # of rows of each subtable. Offset is space between tables
-        expected = 3 * (3 + 4)
-        # 'expand' argument expands the CA dim as slices
-        actual = CrunchCube({}).y_offset(expand=True)
-        self.assertEqual(actual, expected)
-
     @patch('cr.cube.crunch_cube.CrunchCube.is_weighted', False)
     def test_n_unweighted_and_has_no_weight(self):
         unweighted_count = Mock()

tests/unit/test_cube_slice.py

@@ -152,3 +152,35 @@ def test_dim_types(self):
         cube.ndim = 3
         actual = CubeSlice(cube, 0).dim_types
         assert actual == expected
+
+    def test_pruning_2d_labels(self):
+        '''Test that 2D labels are fetched from cr.cube, and pruned.'''
+        cube = Mock()
+        cube.ndim = 2
+        cube.prune_indices.return_value = [
+            np.array([True, False]), np.array([False, False, True]),
+        ]
+        cube.labels.return_value = [
+            [Mock(), 'fake_lbl_1'], ['fake_lbl_2', 'fake_lbl_3', Mock()],
+        ]
+        actual = CubeSlice(cube, 0).labels(prune=True)
+        expected = [['fake_lbl_1'], ['fake_lbl_2', 'fake_lbl_3']]
+        assert actual == expected
+
+    def test_pruning_3d_labels(self):
+        '''Test that 2D labels are fetched from cr.cube, and pruned.'''
+        cube = Mock()
+        cube.ndim = 3
+        cube.prune_indices.return_value = [
+            Mock(),
+            (np.array([True, False]), np.array([False, False, True])),
+            Mock(),
+        ]
+        cube.labels.return_value = [
+            Mock(),
+            [Mock(), 'fake_lbl_1'],
+            ['fake_lbl_2', 'fake_lbl_3', Mock()],
+        ]
+        actual = CubeSlice(cube, 1).labels(prune=True)
+        expected = [['fake_lbl_1'], ['fake_lbl_2', 'fake_lbl_3']]
+        assert actual == expected

tests/unit/test_dimension.py

@@ -180,6 +180,9 @@ def test_dimension_description(self):
         actual = dim.description
         self.assertEqual(actual, expected)
 
+    @patch('cr.cube.dimension.Dimension._elements', [
+        {'id': 1}, {'id': 2}, {'id': 5}, {'id': 4}
+    ])
     @patch('cr.cube.subtotal.Subtotal._all_dim_ids', [1, 2, 4, 5])
     @patch('cr.cube.dimension.Dimension._get_type')
     def test_hs_names_with_bad_data(self, mock_type):
@@ -277,9 +280,11 @@ def test_hs_indices_with_empty_indices(self, mock_type):
         actual = dim.hs_indices
         self.assertEqual(actual, expected)
 
+    @patch('cr.cube.dimension.Dimension.elements')
     @patch('cr.cube.dimension.Dimension._get_type')
-    def test_subtotals(self, mock_type):
+    def test_subtotals(self, mock_type, mock_elements):
         mock_type.return_value = None
+        mock_elements.return_value = [{'id': 1}, {'id': 5}]
         dim_data = {
             'references': {
                 'view': {
@@ -308,43 +313,43 @@ def test_inserted_hs_indices_and_order(self, mock_type):
                         'insertions': [
                             {
                                 u'anchor': u'bottom',
-                                u'args': [],
+                                u'args': [111],
                                 u'function': u'subtotal',
                                 u'name': u'bottoms up one',
                             },
                             {
                                 u'anchor': u'bottom',
-                                u'args': [],
+                                u'args': [222],
                                 u'function': u'subtotal',
                                 u'name': u'bottoms up two',
                             },
                             {
                                 u'anchor': u'bottom',
-                                u'args': [],
+                                u'args': [333],
                                 u'function': u'subtotal',
                                 u'name': u'bottoms up three',
                             },
                             {
                                 u'anchor': u'top',
-                                u'args': [],
+                                u'args': [444],
                                 u'function': u'subtotal',
                                 u'name': u'on top one',
                             },
                             {
                                 u'anchor': u'top',
-                                u'args': [],
+                                u'args': [555],
                                 u'function': u'subtotal',
                                 u'name': u'on top two',
                             },
                             {
                                 u'anchor': 333,
-                                u'args': [],
+                                u'args': [555],
                                 u'function': u'subtotal',
                                 u'name': u'in the middle one',
                             },
                             {
                                 u'anchor': 333,
-                                u'args': [],
+                                u'args': [555],
                                 u'function': u'subtotal',
                                 u'name': u'in the middle two',
                             }

tests/unit/test_subtotal.py

@@ -43,56 +43,80 @@ def test_is_invalid_when_missing_keys(self):
         self.assertEqual(actual, expected)
 
     def test_is_invalid_when_not_subtotal(self):
-        subtotal = Subtotal(self.invalid_subtotal_2, Mock())
+        dim = Mock()
+        dim.elements.return_value = [{'id': 5}, {'id': 4}]
+        subtotal = Subtotal(self.invalid_subtotal_2, dim)
         expected = False
         actual = subtotal.is_valid
         self.assertEqual(actual, expected)
 
     def test_is_valid(self):
-        subtotal = Subtotal(self.valid_subtotal, Mock())
+        dim = Mock()
+        dim.elements.return_value = [{'id': 5}, {'id': 4}]
+        subtotal = Subtotal(self.valid_subtotal, dim)
         expected = True
         actual = subtotal.is_valid
         self.assertEqual(actual, expected)
 
+    def test_is_invalid_when_hs_ids_not_in_dim_elements(self):
+        dim = Mock()
+        dim.elements.return_value = [{'id': 101}, {'id': 102}]
+        subtotal = Subtotal(self.valid_subtotal, dim)
+        expected = False
+        actual = subtotal.is_valid
+        self.assertEqual(actual, expected)
+
     def test_anchor_on_invalid_missing_keys(self):
         subtotal = Subtotal(self.invalid_subtotal_1, Mock())
         expected = None
         actual = subtotal.anchor
         self.assertEqual(actual, expected)
 
     def test_anchor_on_invalid_not_subtotal(self):
-        subtotal = Subtotal(self.invalid_subtotal_2, Mock())
+        dim = Mock()
+        dim.elements.return_value = [{'id': 5}, {'id': 4}]
+        subtotal = Subtotal(self.invalid_subtotal_2, dim)
         expected = None
         actual = subtotal.anchor
         self.assertEqual(actual, expected)
 
     @patch('cr.cube.subtotal.Subtotal._all_dim_ids', [1, 3, 5])
     def test_anchor_on_valid(self):
-        subtotal = Subtotal(self.valid_subtotal, Mock())
+        dim = Mock()
+        dim.elements.return_value = [{'id': 5}, {'id': 4}]
+        subtotal = Subtotal(self.valid_subtotal, dim)
         expected = 5
         actual = subtotal.anchor
         self.assertEqual(actual, expected)
 
     def test_args_on_invalid_1(self):
-        subtotal = Subtotal(self.invalid_subtotal_1, Mock())
+        dim = Mock()
+        dim.elements.return_value = [{'id': 5}, {'id': 4}]
+        subtotal = Subtotal(self.invalid_subtotal_1, dim)
         expected = []
         actual = subtotal.args
         self.assertEqual(actual, expected)
 
     def test_args_on_invalid_2(self):
-        subtotal = Subtotal(self.invalid_subtotal_2, Mock())
+        dim = Mock()
+        dim.elements.return_value = [{'id': 5}, {'id': 4}]
+        subtotal = Subtotal(self.invalid_subtotal_2, dim)
         expected = []
         actual = subtotal.args
         self.assertEqual(actual, expected)
 
     def test_args_on_valid(self):
-        subtotal = Subtotal(self.valid_subtotal, Mock())
+        dim = Mock()
+        dim.elements.return_value = [{'id': 5}, {'id': 4}]
+        subtotal = Subtotal(self.valid_subtotal, dim)
         expected = [5, 4]
         actual = subtotal.args
         self.assertEqual(actual, expected)
 
     def test_anchor_on_uppercased_bottom(self):
-        subtotal = Subtotal(self.valid_subtotal_anchor_bottom, Mock())
+        dim = Mock()
+        dim.elements.return_value = [{'id': 5}, {'id': 4}]
+        subtotal = Subtotal(self.valid_subtotal_anchor_bottom, dim)
         anchor = subtotal.anchor
         assert anchor == 'bottom'