Advertisement
Uno-Dan

pTree

Jun 20th, 2020
1,455
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 7.30 KB | None | 0 0
  1. from enum import IntEnum
  2. from collections import deque
  3. import config
  4.  
  5. const = IntEnum('Constants', 'START END', start=-1)
  6.  
  7.  
  8. class Base:
  9.     def __init__(self, data=None, **kwargs):
  10.         self.id = kwargs.get('id')
  11.         self.name = kwargs.get('name')
  12.         self.parent = kwargs.get('parent')
  13.         self.columns = kwargs.get('columns')
  14.  
  15.         if data:
  16.             self.id = data.get('id')
  17.             self.name = data.get('name')
  18.             self.parent = data.get('parent')
  19.             self.columns = data.get('columns')
  20.  
  21.     @property
  22.     def tree(self):
  23.         item = self
  24.         while not isinstance(item, Tree):
  25.             item = item.parent
  26.         return item
  27.  
  28.     def get(self, column):
  29.         if not column:
  30.             return self.name
  31.         elif self and len(self.columns) > column >= 1:
  32.             return self.columns[column-1]
  33.  
  34.     def set(self, column, value):
  35.         if column < 0:
  36.             return
  37.         if column > len(self.columns):
  38.             self.columns.append(value)
  39.         else:
  40.             self.columns[column-1] = value
  41.  
  42.     def is_node(self, item=None):
  43.         if not item:
  44.             item = self
  45.         return True if isinstance(item, Node) else False
  46.  
  47.  
  48. class Leaf(Base):
  49.     def __init__(self, data=None,  **kwargs):
  50.         super().__init__(data, **kwargs)
  51.  
  52.         if self.id is None and self.parent:
  53.             self.id = self.tree.next_id()
  54.         elif self.parent:
  55.             for node in self.parent:
  56.                 if node.id == self.id:
  57.                     self.id = self.tree.next_id()
  58.                     return
  59.         else:
  60.             self.id = self.tree.next_id()
  61.  
  62.  
  63. class Node(Base, deque):
  64.     def __init__(self, data=None, **kwargs):
  65.         Base.__init__(self, data, **kwargs)
  66.         deque.__init__(self)
  67.  
  68.         if self.id is None and self.parent:
  69.             self.id = self.tree.next_id()
  70.         elif self.parent:
  71.             for node in self.parent:
  72.                 if node.id == self.id:
  73.                     self.id = self.tree.next_id()
  74.                     return
  75.             self.id = self.tree.next_id()
  76.         elif not isinstance(self, Tree):
  77.             self.id = self.tree.next_id()
  78.  
  79.     @property
  80.     def children(self):
  81.         return self
  82.  
  83.     def show(self, parent=None, indent=3):
  84.         def walk(_parent, level=0):
  85.             for idx, _node in enumerate(_parent):
  86.                 pad = '' if not level else ' ' * (indent * level)
  87.                 print(f' {_node.id}:{pad} {_node.name}, {_node.columns}')
  88.                 if _parent.is_node(_node):
  89.                     walk(_node, level+1)
  90.  
  91.         print('-----------------------------------------------------')
  92.         print(f'ID: Name, Columns: {str(self.tree.headings)}')
  93.         print('-----------------------------------------------------')
  94.         walk(parent if parent else self)
  95.         print('-----------------------------------------------------')
  96.  
  97.     def query(self, query):
  98.         if isinstance(query, int):
  99.             return self.query_by_id(query)
  100.         elif isinstance(query, str):
  101.             return self.query_by_name(query)
  102.  
  103.     def populate(self, data, **kwargs):
  104.         self.tree.populate(data, parent=self)
  105.  
  106.     def get_cell(self, row, column):
  107.         item = self.query(row)
  108.         if not item:
  109.             return
  110.         if column:
  111.             return item.get(column)
  112.         else:
  113.             return self.name
  114.  
  115.     def set_cell(self, row, column, value):
  116.         item = self.query(row)
  117.         if column:
  118.             item.set(column, value)
  119.         else:
  120.             item.name = value
  121.  
  122.     def query_by_id(self, _id):
  123.         def walk(_item):
  124.             if _item.id == _id:
  125.                 return _item
  126.             if not _item.is_node():
  127.                 return
  128.  
  129.             for child in _item:
  130.                 if child.is_node():
  131.                     _result = walk(child)
  132.                     if _result:
  133.                         return _result
  134.                 elif child.id == _id:
  135.                     return child
  136.  
  137.         for item in self:
  138.             result = walk(item)
  139.             if result:
  140.                 return result
  141.  
  142.     def query_by_name(self, name):
  143.         def walk(_item):
  144.             if _item.name == name:
  145.                 return _item
  146.             if not _item.is_node():
  147.                 return
  148.  
  149.             for child in _item:
  150.                 if child.is_node():
  151.                     _result = walk(child)
  152.                     if _result:
  153.                         return _result
  154.                 elif child.name == name:
  155.                     return child
  156.  
  157.         for item in self:
  158.             result = walk(item)
  159.             if result:
  160.                 return result
  161.  
  162.  
  163. class Tree(Node):
  164.     def __init__(self, data=None):
  165.         self.size = 0
  166.         self.headings = ['Type', 'Size', 'Parent']
  167.         super().__init__()
  168.  
  169.         self.name = '.'
  170.         self.parent = None
  171.         self.populate(data)
  172.  
  173.     def next_id(self):
  174.         self.size += 1
  175.         return self.size
  176.  
  177.     def reindex(self, start=0):
  178.         def walk(_parent, level=0):
  179.             for _node in _parent:
  180.                 _node.id = self.next_id()
  181.                 if _parent.is_node(_node):
  182.                     walk(_node, level+1)
  183.  
  184.         self.size = start
  185.         walk(self)
  186.  
  187.     def populate(self, data, **kwargs):
  188.         if not data:
  189.             return
  190.  
  191.         def walk(_parent, _item):
  192.             if 'children' in _item:
  193.                 _item['columns'] = ['Node', '0 items', _parent.name]
  194.                 item = Node(parent=_parent, **_item)
  195.                 if 'children' in _item and len(_item['children']):
  196.                     for node in _item['children']:
  197.                         walk(item, node)
  198.             else:
  199.                 _item['columns'] = ['Leaf', '0 Kb', _parent.name]
  200.                 item = Leaf(parent=_parent, **_item)
  201.             _parent.append(item)
  202.  
  203.         parent = kwargs.get('parent', self)
  204.         if isinstance(data, list):
  205.             for cfg in data:
  206.                 walk(parent, cfg)
  207.  
  208.  
  209. def main():
  210.     t = Tree()
  211.     # Create node with kwargs.
  212.     node = Node(parent=t, name='Test 4', columns=['Node', '0 items', t.name])
  213.     t.append(node)
  214.  
  215.     leaf = Leaf(parent=t, name='Test 1', columns=['Leaf', '0 Kb', t.name])
  216.     t.append(leaf)
  217.  
  218.     # Create node with dict.
  219.     node = Node({'parent': t, 'name': 'Test 2', 'columns': ['Node', '0 items', t.name]})
  220.     t.append(node)
  221.  
  222.     node = Node(parent=t, name='Test 3', columns=['Node', '0 items', t.name])
  223.     t.append(node)
  224.  
  225.     leaf = Leaf(parent=node, name='Leaf 3-1', columns=['Leaf', '0 Kb', node.name])
  226.     node.append(leaf)
  227.  
  228.     t.show()
  229.  
  230.     cfg = [{
  231.         'name': 'Node 1a',
  232.         'children': [
  233.             {
  234.                 'name': 'Leaf 1a',
  235.             },
  236.             {
  237.                 'name': 'Leaf 2a',
  238.             },
  239.             {
  240.                 'name': 'Sub Node 1a',
  241.                 'children': [
  242.                     {
  243.                         'name': 'Sub Leaf 1a',
  244.                     },
  245.                 ]
  246.             },
  247.         ],
  248.     }]
  249.  
  250.     t = Tree(cfg)
  251.  
  252.     node = t[0]
  253.     node.set(1, 'Test')
  254.     t.show()
  255.  
  256.     t.set_cell(5, 0, 'New Name')
  257.     t.show()
  258.  
  259.     t.query(4).populate(cfg)
  260.     t.show()
  261.  
  262.  
  263. if __name__ == '__main__':
  264.     main()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement