tree.py

# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
from treelib import Node, Tree
import sys

# ------------------------------------------------------------------------------------------
#
# Read the CSV file exported from Confluence.py
#
# ------------------------------------------------------------------------------------------

df = pd.read_csv("confluence.csv") # Make sure this is in the same directy as this file

# ------------------------------------------------------------------------------------------
#
# Create a tree using treelib and make a base node named 'Confluence'
#
# ------------------------------------------------------------------------------------------

tree = Tree() # Create an empty tree object to populate
tree.create_node("Confluence", "confluence") # Create the first node; there can only be one top level node

# ------------------------------------------------------------------------------------------
#
# Create empty dataframes for later use
#
# ------------------------------------------------------------------------------------------

newDF = pd.DataFrame(columns=['id', 'space', 'parent', 'name', 'url'])
secondTierDF = pd.DataFrame(columns=['id', 'space', 'parent', 'name', 'url'])
thirdTierDF = pd.DataFrame(columns=['id', 'space', 'parent', 'name', 'url'])
fourthTierDF = pd.DataFrame(columns=['id', 'space', 'parent', 'name', 'url'])

# ------------------------------------------------------------------------------------------
#
# Create first level heirarchy from spaces
#
# ------------------------------------------------------------------------------------------

# Create first level of tree made up of confluence spaces
for index, row in df.iterrows():
    if row['parent'] == 0:
        # Clean up dataframe to remove tildes from space names (generated by user pages)
        if not row['space'].find("~"):
            row['space'] = row['space'][1:]
            spaceName = row[3]
            tree.create_node(spaceName, row['id'], parent="confluence")
            newDF = newDF.append(row)
            df.drop(index) #drop all personal pages
        else:
            tree.create_node(row['space'], row['id'], parent="confluence")
            newDF = newDF.append(row)

# Creating a copy of the original page DF to be used later on
df2 = df

# ------------------------------------------------------------------------------------------
#
# Loop through 3 additional levels of page hierarchies for each space
#
# ------------------------------------------------------------------------------------------

# Create the second level heirarchy by comparing every page ID to every parent ID
for index, row in newDF.iterrows():
    parentID = row['id']
    for index2, confRow in df.iterrows():
        childID = confRow['parent']
        if parentID == childID:
            pageName = confRow[3]
            print("TIER 1", pageName)
            tree.create_node(pageName, confRow[0], parent=parentID)
            secondTierDF = secondTierDF.append(confRow)
            df.drop(df.index[index2])


# Create third level heirarchy
for index, row in secondTierDF.iterrows():
    parentID2 = row['id']
    for index2, secondTierRow in df2.iterrows():
        childID2 = secondTierRow['parent']
        if parentID2 == childID2:
            pageName2 = secondTierRow[3]
            print("TIER 2", pageName2)
            tree.create_node(pageName2, secondTierRow[0], parent=parentID2)
            thirdTierDF = thirdTierDF.append(secondTierRow)

# Create fourth level heirarchy
for index, row in thirdTierDF.iterrows():
    parentID3 = row['id']
    for index3, thirdTierRow in df2.iterrows():
        childID3 = thirdTierRow['parent']
        if parentID3 == childID3:
            pageName3 = thirdTierRow[3]
            print("TIER 3", pageName3)
            tree.create_node(pageName3, thirdTierRow[0], parent=parentID3)

# Create fifth level heirarchy
for index, row in fourthTierDF.iterrows():
    parentID4 = row['id']
    for index4, fourthTierRow in df2.iterrows():
        childID4 = fourthTierRow['parent']
        if parentID4 == childID4:
            pageName4 = fourthTierRow[3]
            print("TIER 4!", pageName4)
            tree.create_node(pageName4, fourthTierRow[0], parent=parentID4)

# ------------------------------------------------------------------------------------------
#
# Print the tree generated by treelib
#
# ------------------------------------------------------------------------------------------

print(tree.show()) # Prints the tree in the console