prodir/tree_structurer/cpp/tree_structurer.cpp

#include "tree_structurer.hpp"
#include <algorithm>
#include <filesystem>

namespace fs = std::filesystem;

bool TreeStructurer::should_ignore_dir(const std::string& dirname) {
    static const std::vector<std::string> ignore_list = {
        "build", "venv", "myenv", "dist", "node_modules", "CMakeFiles",
        ".git", ".idea", ".vscode", "__pycache__", "**pycache**"
    };

    // Check for __main__.py or __init__.py files
    if (!dirname.empty() && (dirname == "__main__.py" || dirname == "__init__.py")) {
        return false;  // Do not ignore these files
    }

    if (std::find(ignore_list.begin(), ignore_list.end(), dirname) != ignore_list.end()) {
        return true;
    }

    if (!dirname.empty()) {
        if (dirname[0] == '.' || dirname[0] == '_') {
            return true;
        }
        if (dirname.find("__") == 0 && dirname.find("__", 2) != std::string::npos) {
            return true;
        }
    }

    return false;
}

bool TreeStructurer::should_ignore_file(const std::string& filename) {
    static const std::vector<std::string> ignore_extensions = {
        ".pyc", ".pyo", ".pyd", ".so", ".dll", ".dylib",
        ".o", ".obj", ".a", ".lib"
    };

    // Check for __main__.py or __init__.py files
    if (!filename.empty() && (filename == "__main__.py" || filename == "__init__.py")) {
        return false;  // Do not ignore these files
    }

    if (!filename.empty() && (filename[0] == '.' || filename[0] == '_')) {
        return true;
    }

    fs::path path(filename);
    std::string ext = path.extension().string();
    if (std::find(ignore_extensions.begin(), ignore_extensions.end(), ext) != ignore_extensions.end()) {
        return true;
    }

    return false;
}

// Add the missing get_relative_path implementation
std::string TreeStructurer::get_relative_path(const fs::path& path, const fs::path& base) {
    fs::path rel = fs::relative(path, base);
    return rel.string();
}

std::vector<fs::path> TreeStructurer::get_filtered_paths(const fs::path& start) {
    std::vector<fs::path> paths;
    fs::directory_options options = fs::directory_options::skip_permission_denied;

    try {
        if (!fs::exists(start)) {
            throw std::runtime_error("Directory does not exist: " + start.string());
        }
        
        if (!fs::is_directory(start)) {
            throw std::runtime_error("Path is not a directory: " + start.string());
        }

        paths.push_back(start);

        // Check if directory is empty
        bool is_empty = fs::directory_iterator(start) == fs::directory_iterator();
        if (is_empty) {
            throw std::runtime_error("Directory is empty: " + start.string());
        }

        for (const auto& entry : fs::recursive_directory_iterator(start, options)) {
            const auto& path = entry.path();

            bool should_skip = false;
            for (const auto& component : path) {
                if (should_ignore_dir(component.string())) {
                    should_skip = true;
                    break;
                }
            }
            if (should_skip) continue;

            if (entry.is_directory()) {
                if (!should_ignore_dir(path.filename().string())) {
                    paths.push_back(path);
                }
            } else {
                if (!should_ignore_file(path.filename().string())) {
                    paths.push_back(path);
                }
            }
        }
    } catch (const fs::filesystem_error& e) {
        throw std::runtime_error("Error accessing path: " + std::string(e.what()));
    }

    std::sort(paths.begin(), paths.end());
    return paths;
}

std::vector<std::string> TreeStructurer::get_directory_structure(const std::string& startpath) {
    std::vector<std::string> result;
    
    // Normalize the input path by removing ./ or .\ prefix if present
    std::string normalized_path = startpath;
    if (startpath.substr(0, 2) == ".\\" || startpath.substr(0, 2) == "./") {
        normalized_path = startpath.substr(2);
    }
    
    fs::path start = normalized_path.empty() ? fs::current_path() : fs::path(normalized_path);

    try {
        auto paths = get_filtered_paths(start);
        if (paths.empty()) {
            throw std::runtime_error("No valid files or directories found in: " + start.string());
        }
        std::vector<bool> is_last_at_level(256, false);
        
        for (size_t i = 1; i < paths.size(); ++i) {
            const auto& path = paths[i];
            std::string rel_path = get_relative_path(path, start);
            
            int level = std::count(rel_path.begin(), rel_path.end(), fs::path::preferred_separator);

            bool is_last = true;
            for (size_t j = i + 1; j < paths.size(); ++j) {
                std::string next_rel_path = get_relative_path(paths[j], start);
                int next_level = std::count(next_rel_path.begin(), next_rel_path.end(), fs::path::preferred_separator);
                if (next_level == level) {
                    is_last = false;
                    break;
                }
                if (next_level < level) {
                    break;
                }
            }

            is_last_at_level[level] = is_last;

            std::string line;
            for (int j = 0; j < level; ++j) {
                if (j == level - 1) {
                    line += is_last ? "└── " : "├── ";
                } else {
                    line += is_last_at_level[j] ? "    " : "│   ";
                }
            }

            line += path.filename().string();
            if (fs::is_directory(path)) {
                line += "/";
            }
            result.push_back(line);
        }
    } catch (const fs::filesystem_error& e) {
        throw std::runtime_error("Failed to access directory: " + std::string(e.what()));
    }

    return result;
}