python-layout/pillow_layout/internal/flexlayouter.py

from ..enums import FlexDirection, FlexWrap, FlexJustify, FlexAlignItems, FlexAlignContent
from .helpers import min_with_none, max_with_none


class FlexLineItem:
    def __init__(self, content, max_width, max_height, direction):
        self.__content = content
        self.__max_width = max_width
        self.__max_height = max_height
        self.__direction = direction
        self.__target_main_size = None
        self.__main_size = None
        self.__cross_size = None
        self.__main_pos = None
        self.__cross_pos = None
        self.__frozen = False

    def content(self):
        return self.__content

    def base_size(self):
        if self.__direction in (FlexDirection.ROW, FlexDirection.ROW_REVERSE):
            if self.__content.width is not None:
                return self.__content.width
            else:
                return self.__content.get_ideal_outer_dimensions(available_width=self.__max_width,
                                                                 available_height=self.__max_height)[0]
        else:
            if self.__content.height is not None:
                return self.__content.height
            else:
                return self.__content.get_ideal_outer_dimensions(available_width=self.__max_width,
                                                                 available_height=self.__max_height)[1]

    def min_main_size(self):
        if self.__direction in (FlexDirection.ROW, FlexDirection.ROW_REVERSE):
            return self.__content.get_min_outer_width()
        else:
            return self.__content.get_min_outer_height()

    def min_cross_size(self):
        if self.__direction in (FlexDirection.ROW, FlexDirection.ROW_REVERSE):
            return self.__content.get_min_outer_height()
        else:
            return self.__content.get_min_outer_width()

    def max_main_size(self):
        if self.__direction in (FlexDirection.ROW, FlexDirection.ROW_REVERSE):
            return self.__content.max_width
        else:
            return self.__content.max_height

    def max_cross_size(self):
        if self.__direction in (FlexDirection.ROW, FlexDirection.ROW_REVERSE):
            return self.__content.max_height
        else:
            return self.__content.max_width

    def hypothetical_main_size(self):
        return min_with_none(self.max_main_size(), max_with_none(self.min_main_size(), self.base_size()))

    def target_main_size(self):
        return self.__target_main_size

    def main_size(self):
        return self.__main_size

    def hypothetical_cross_size(self):
        if self.__direction in (FlexDirection.ROW, FlexDirection.ROW_REVERSE):
            return self.__content.get_ideal_outer_dimensions(available_width=self.__main_size)[1]
        else:
            return self.__content.get_ideal_outer_dimensions(available_height=self.__main_size)[0]

    def cross_size(self):
        return self.__cross_size

    def main_pos(self):
        return self.__main_pos

    def cross_pos(self):
        return self.__cross_pos

    def grow_factor(self):
        return self.__content.flex_grow

    def shrink_factor(self):
        return self.__content.flex_shrink

    def scaled_shrink_factor(self):
        return self.shrink_factor() * self.base_size()

    def set_target_main_size(self, size):
        self.__target_main_size = size

    def adjust_target_main_size(self, adjust):
        self.__target_main_size += adjust

    def set_cross_size(self, size):
        size = min_with_none(max_with_none(size, self.min_cross_size()), self.max_cross_size())
        self.__cross_size = size

    def set_main_pos(self, pos):
        self.__main_pos = pos

    def set_cross_pos(self, pos):
        self.__cross_pos = pos

    def frozen(self):
        return self.__frozen

    def freeze(self):
        self.__frozen = True
        self.__main_size = self.__clamped_target_main_size()

    def __clamped_target_main_size(self):
        return min_with_none(max_with_none(self.__target_main_size, self.min_main_size()), self.max_main_size())

    def violation_size(self):
        return self.__clamped_target_main_size() - self.__target_main_size


class FlexLine():
    def __init__(self, items, gap, direction, justify, align):
        self.__items = items
        self.__gap = gap
        self.__direction = direction
        self.__justify = justify
        self.__align = align
        self.__cross_size = None
        self.__cross_pos = None

    def items(self):
        return self.__items[:]

    def cross_size(self):
        return self.__cross_size

    def set_cross_pos(self, pos):
        self.__cross_pos = pos

    def cross_pos(self):
        return self.__cross_pos

    def compute_main_sizes(self, available_main_space):
        sum_main_sizes = sum([item.hypothetical_main_size() for item in self.__items]) + self.__gap * (
                len(self.__items) - 1)
        if sum_main_sizes > available_main_space:
            factor_type = 'shrink'
        elif sum_main_sizes < available_main_space:
            factor_type = 'grow'
        else:
            for item in self.__items:
                item.set_target_main_size(item.hypothetical_main_size())
                item.freeze()
            return
        for item in self.__items:
            flex_factor = item.grow_factor() if factor_type == 'grow' else item.shrink_factor()
            if (flex_factor == 0 or
                    (factor_type == 'grow' and item.base_size() > item.hypothetical_main_size()) or
                    (factor_type == 'shrink' and item.base_size() < item.hypothetical_main_size())):
                item.set_target_main_size(item.hypothetical_main_size())
                item.freeze()
        initial_used_space = self.__gap * (len(self.__items) - 1)
        for item in self.__items:
            if item.frozen():
                initial_used_space += item.target_main_size()
            else:
                initial_used_space += item.base_size()
        initial_free_space = available_main_space - initial_used_space
        while True:
            unfrozen_items = [
                item for item in self.__items if not item.frozen()]
            if len(unfrozen_items) == 0:
                break
            remaining_used_space = self.__gap * (len(self.__items) - 1)
            for item in self.__items:
                if item.frozen():
                    remaining_used_space += item.main_size()
                else:
                    remaining_used_space += item.base_size()
            remaining_free_space = available_main_space - remaining_used_space
            flex_sum = sum([item.grow_factor() if factor_type == 'grow' else item.shrink_factor()
                            for item in unfrozen_items])
            if flex_sum < 1:
                if abs(initial_free_space * flex_sum) < abs(remaining_free_space):
                    remaining_free_space = initial_free_space * flex_sum
            if remaining_free_space != 0:
                for item in unfrozen_items:
                    item.set_target_main_size(item.base_size())
                if factor_type == 'grow':
                    remaining = remaining_free_space
                    items = sorted(
                        unfrozen_items, key=lambda x: x.grow_factor())
                    while remaining > 0 and len(items) > 0:
                        ratios = [
                            item.grow_factor() / flex_sum for item in items]
                        used = 0
                        for i in range(len(items)):
                            size = int(ratios[i] * remaining)
                            used += size
                            items[i].adjust_target_main_size(size)
                        remaining = remaining - used
                        items = items[:-1]
                else:
                    remaining = remaining_free_space
                    items = sorted(
                        unfrozen_items, key=lambda x: x.scaled_shrink_factor())
                    flex_shrink_sum = sum(
                        [item.scaled_shrink_factor() for item in unfrozen_items])
                    while remaining < 0 and len(items) > 0:
                        ratios = [item.scaled_shrink_factor(
                        ) / flex_shrink_sum for item in items]
                        used = 0
                        for i in range(len(items)):
                            size = int(ratios[i] * remaining)
                            used += size
                            items[i].adjust_target_main_size(size)
                        remaining = remaining - used
                        items = items[:-1]
            sum_violations = sum([item.violation_size()
                                  for item in unfrozen_items])
            if sum_violations > 0:
                for item in unfrozen_items:
                    if item.violation_size() > 0:
                        item.freeze()
            elif sum_violations < 0:
                for item in unfrozen_items:
                    if item.violation_size() < 0:
                        item.freeze()
            else:
                for item in unfrozen_items:
                    item.freeze()
        for item in self.__items:
            item.freeze()

    def compute_cross_sizes(self, min_cross_space, max_cross_space):
        line_size = 0
        for item in self.__items:
            size = item.hypothetical_cross_size()
            item.set_cross_size(size)
            line_size = max(line_size, size)
        if min_cross_space is not None:
            line_size = max(line_size, min_cross_space)
        if max_cross_space is not None:
            line_size = min(line_size, max_cross_space)
        self.__cross_size = line_size

    def set_cross_size(self, size):
        self.__cross_size = size

    def adjust_cross_size(self, adjust):
        self.__cross_size += adjust

    def arrange(self, available_main_space):
        if self.__align == FlexAlignItems.START:
            for item in self.__items:
                item.set_cross_pos(0)
        elif self.__align == FlexAlignItems.END:
            for item in self.__items:
                item.set_cross_pos(self.__cross_size - item.cross_size())
        elif self.__align == FlexAlignItems.CENTER:
            for item in self.__items:
                item.set_cross_pos((self.__cross_size - item.cross_size()) // 2)
        elif self.__align == FlexAlignItems.STRETCH:
            for item in self.__items:
                item.set_cross_size(self.__cross_size)
                item.set_cross_pos(0)

        num_items = len(self.__items)
        num_inner_gaps = num_items - 1
        num_gaps = num_inner_gaps + 2
        gaps = [0] + [self.__gap] * num_inner_gaps + [0]
        remaining_space = available_main_space - self.__gap * num_inner_gaps - sum(
            [item.main_size() for item in self.__items])
        justify = self.__justify
        if remaining_space < 0 and justify == FlexJustify.SPACE_BETWEEN:
            justify = FlexJustify.START
        elif remaining_space < 0 and justify in (FlexJustify.SPACE_EVENLY, FlexJustify.SPACE_AROUND):
            justify = FlexJustify.CENTER
        if justify == FlexJustify.START:
            gaps[-1] += remaining_space
        elif justify == FlexJustify.END:
            gaps[0] += remaining_space
        elif justify == FlexJustify.CENTER:
            gaps[0] += remaining_space // 2
            gaps[-1] += remaining_space - (remaining_space // 2)
        elif justify == FlexJustify.SPACE_BETWEEN:
            space_parts = [remaining_space // num_inner_gaps + (1 if x < remaining_space % num_inner_gaps else 0) for x
                           in range(num_inner_gaps)]
            for i in range(num_inner_gaps):
                gaps[i + 1] += space_parts[i]
        elif justify == FlexJustify.SPACE_AROUND:
            space_parts = [remaining_space // num_items + (1 if x < remaining_space % num_items else 0) for x in
                           range(num_items)]
            for i in range(num_items):
                gaps[i] += space_parts[i] // 2
                gaps[i + 1] += space_parts[i] - (space_parts[i] // 2)
        elif justify == FlexJustify.SPACE_EVENLY:
            space_parts = [remaining_space // num_gaps + (1 if x < remaining_space % num_gaps else 0) for x in
                           range(num_gaps)]
            for i in range(num_gaps):
                gaps[i] += space_parts[i]
        else:
            raise Exception('invalid flex justify')

        if self.__direction in (FlexDirection.ROW, FlexDirection.COLUMN):
            main_pos = 0
            for i in range(num_items):
                item = self.__items[i]
                space_before = gaps[i]
                main_pos += space_before
                item.set_main_pos(main_pos)
                main_pos += item.main_size()
        elif self.__direction in (FlexDirection.ROW_REVERSE, FlexDirection.COLUMN_REVERSE):
            main_pos = available_main_space
            for i in range(num_items):
                item = self.__items[i]
                space_before = gaps[i]
                main_pos -= space_before
                main_pos -= item.main_size()
                item.set_main_pos(main_pos)
        else:
            raise Exception('invalid flex direction')


class FlexLayouter():
    def __init__(self, items, gap, direction, wrap, justify, align_content, align_items):
        self.__items = items
        self.__gap = gap
        self.__direction = direction
        self.__wrap = wrap
        self.__justify = justify
        self.__align_content = align_content
        self.__align_items = align_items

    def arrange(self, min_width=None, min_height=None, max_width=None, max_height=None, consume_space=True):
        if self.__direction in (FlexDirection.ROW, FlexDirection.ROW_REVERSE):
            available_main_space = max_width
            available_cross_space = max_height
            min_cross_space = min_height
        else:
            available_main_space = max_height
            available_cross_space = max_width
            min_cross_space = min_width

        items = [FlexLineItem(i, max_width, max_height, self.__direction) for i in self.__items]

        if self.__wrap == FlexWrap.NO_WRAP or available_main_space is None:
            lines = [items]
            if available_main_space is None:
                available_main_space = self.__gap * (len(items) - 1) + sum(
                    [item.hypothetical_main_size() for item in items])
        else:
            lines = []
            current_line = []
            current_line_length = 0
            for item in items:
                size = item.hypothetical_main_size()
                new_line_length = current_line_length + size
                if len(current_line) > 0:  # not the first item in the line, so it needs a gap before it
                    new_line_length += self.__gap
                if new_line_length <= available_main_space:  # item fits into the current line
                    current_line.append(item)
                    current_line_length = new_line_length
                else:  # item does not fit into the current line
                    # current line already has items, so it gets completed first
                    if len(current_line) > 0:
                        lines.append(current_line)
                    current_line = [item]
                    current_line_length = size
            if len(current_line) > 0:
                lines.append(current_line)

        lines = [FlexLine(items, self.__gap, self.__direction,
                          self.__justify, self.__align_items) for items in lines]
        num_lines = len(lines)

        for line in lines:
            line.compute_main_sizes(available_main_space)
        for line in lines:
            if num_lines > 1:
                line.compute_cross_sizes(None, None)
            else:
                line.compute_cross_sizes(min_cross_space, available_cross_space)
                if available_cross_space is not None and consume_space:
                    line.set_cross_size(available_cross_space)
        lines_sum_cross_space = self.__gap * (len(lines) - 1) + sum([line.cross_size() for line in lines])
        if (consume_space and
                self.__align_content == FlexAlignContent.STRETCH and
                available_cross_space is not None and
                lines_sum_cross_space < available_cross_space):
            remaining_cross_space = available_cross_space - lines_sum_cross_space
            distributed_cross_space = [remaining_cross_space // num_lines + (
                1 if x < remaining_cross_space % num_lines else 0) for x in range(num_lines)]
            for i in range(num_lines):
                lines[i].adjust_cross_size(distributed_cross_space[i])
        for line in lines:
            line.arrange(available_main_space)
        if len(lines) == 1:
            lines[0].set_cross_pos(0)
        else:
            align_content = self.__align_content
            num_lines = len(lines)
            num_inner_gaps = num_lines - 1
            num_gaps = num_inner_gaps + 2
            used_space = self.__gap * num_inner_gaps + sum([line.cross_size() for line in lines])
            gaps = [0] + [self.__gap] * num_inner_gaps + [0]
            if available_cross_space is None:
                # cannot do any other alignments without a definite cross size
                align_content = FlexAlignContent.START
                available_cross_space = used_space
            remaining_space = available_cross_space - used_space
            # stretch was already handled before, and behaves like start otherwise
            if align_content == FlexAlignContent.START or align_content == FlexAlignContent.STRETCH:
                gaps[-1] += remaining_space
            elif align_content == FlexAlignContent.END:
                gaps[0] += remaining_space
            elif align_content == FlexAlignContent.CENTER:
                gaps[0] += remaining_space // 2
                gaps[-1] += remaining_space - (remaining_space // 2)
            elif align_content == FlexAlignContent.SPACE_BETWEEN:
                space_parts = [remaining_space // num_inner_gaps + (1 if x < remaining_space % num_inner_gaps else 0)
                               for x in range(num_inner_gaps)]
                for i in range(num_inner_gaps):
                    gaps[i + 1] += space_parts[i]
            elif align_content == FlexAlignContent.SPACE_AROUND:
                space_parts = [remaining_space // num_lines + (1 if x < remaining_space % num_lines else 0) for x in
                               range(num_lines)]
                for i in range(num_lines):
                    gaps[i] += space_parts[i] // 2
                    gaps[i + 1] += space_parts[i] - (space_parts[i] // 2)
            elif align_content == FlexAlignContent.SPACE_EVENLY:
                space_parts = [remaining_space // num_gaps + (1 if x < remaining_space % num_gaps else 0) for x in
                               range(num_gaps)]
                for i in range(num_gaps):
                    gaps[i] += space_parts[i]
            else:
                raise Exception('invalid flex align content')

            if self.__wrap == FlexWrap.WRAP:
                cross_pos = 0
                for i in range(num_lines):
                    line = lines[i]
                    space_before = gaps[i]
                    cross_pos += space_before
                    line.set_cross_pos(cross_pos)
                    cross_pos += line.cross_size()
            elif self.__wrap == FlexWrap.WRAP_REVERSE:
                cross_pos = available_cross_space
                for i in range(num_lines):
                    line = lines[i]
                    space_before = gaps[i]
                    cross_pos -= space_before
                    cross_pos -= line.cross_size()
                    line.set_cross_pos(cross_pos)
            else:
                raise Exception('invalid flex wrap')

        items_with_rects = []
        horizontal = self.__direction in (FlexDirection.ROW, FlexDirection.ROW_REVERSE)
        for line in lines:
            if horizontal:
                line_x = 0
                line_y = line.cross_pos()
            else:
                line_x = line.cross_pos()
                line_y = 0
            for item in line.items():
                if horizontal:
                    item_x = item.main_pos()
                    item_y = item.cross_pos()
                    item_width = item.main_size()
                    item_height = item.cross_size()
                else:
                    item_x = item.cross_pos()
                    item_y = item.main_pos()
                    item_width = item.cross_size()
                    item_height = item.main_size()
                x1 = line_x + item_x
                y1 = line_y + item_y
                x2 = x1 + item_width - 1
                y2 = y1 + item_height - 1
                items_with_rects.append(((x1, y1, x2, y2), item.content()))
        return items_with_rects

    def get_dimensions(self, min_width=None, min_height=None, max_width=None, max_height=None):
        contents = self.arrange(min_width, min_height, max_width, max_height, consume_space=False)
        max_x = 0
        max_y = 0
        for rect, _ in contents:
            _, _, x2, y2 = rect
            max_x = max(max_x, x2)
            max_y = max(max_y, y2)
        return max_x, max_y