# This file is part of Tryton.  The COPYRIGHT file at the top level of
# this repository contains the full copyright notices and license terms.

from trytond.pool import PoolMeta

__all__ = ['CostDistribution', 'DistributionRule']

__metaclass__ = PoolMeta


class DistributionRule:

    __name__ = 'cost.manage.distribution.rule'

    @classmethod
    def _get_cost_source(cls):
        """Return list of Model names for Cost source Reference"""
        sources = super(DistributionRule, cls)._get_cost_source()
        sources.append('production')
        return sources


class CostDistribution:

    __name__ = 'cost.manage.cost.distribution'

    @classmethod
    def _get_cost_center(cls):
        """Return list of Model names for Cost center Reference"""
        cost_centers = super(CostDistribution, cls)._get_cost_center()
        cost_centers.append('production')
        return cost_centers