R/functions.r

In pkgGraphR: Graph the Relationship Between Functions in an R Package

#' Collect all functions in a package or directory
#'
#' @description collect all the functions defined in an R program, directory, or file
#' @name collectFunNames
#'
#' @param x A character string specifying the path to an R package, directory, or file
#'
#' @return A named list of function assignments in each `.R` file in `x`
#'
#' @examples
#' system.file("extdata", package = "pkgGraphR") |>
#'     collectFunNames()
#'
#' @export
collectFunNames <- function(x){
    stopifnot(is.character(x), length(x) == 1)
    x <- normalizePath(x)
    if(dir.exists(x)){
        Rpath <- list.dirs(x)[grep("^R$", list.dirs(x, full.names = FALSE))]
        if(length(Rpath) == 0){
            Rpath <- grep("/\\.", list.dirs(x), value = TRUE, invert = TRUE)
            if(length(Rpath) == 0){
                stop("'x' is an empty directory")
            }
        }
        Rscripts <- grep("\\.r$|\\.R$", list.files(Rpath, full.names = TRUE), value = TRUE)
        if(length(Rscripts) == 0){
            stop("'x' does not contain any .R or .r files.")
        }
    } else if(file.exists(x)){
        Rscripts <- grep(".R$|.r$", x, value = TRUE)
        if(length(Rscripts) == 0){
            stop("'x' is a file, but not a .R or .r file.")
        }
    } else {
        stop("'x' must be a valid file or directory")
    }

    allFuns <- sapply(Rscripts, \(x){
        pR <- parse(x, keep.source = TRUE)
        pRD <- utils::getParseData(pR) |> dplyr::filter(token != "COMMENT", terminal != FALSE)
        findFunNames(pRD)
    }, simplify = FALSE, USE.NAMES = TRUE) |> purrr::keep(\(x) nrow(x) > 0)
    # outside of above lapply to avoid empty files problem
    allFuns <- lapply(allFuns, \(x){
        dplyr::filter(x, !grepl("[]$:<[`-]", x = name))
    })
    return(allFuns)
}

#' Build a graph of an R package or directory
#'
#' @description Generates the Nodes and Edges of a set of functions in an R package or directory
#' @name buildPackageGraph
#'
#' @param x A character string specifying the path to an R package or directory
#' @param unique.edges Logical indicating whether there should be only a single
#' edge between nodes.
#' DEFAULT: TRUE
#' @param only.connected Logical indicating whether unconnected nodes should be
#' removed from the graph.
#' DEFAULT: FALSE
#'
#' @return A named list of length 2 containing a character vector of nodes and a
#' data.frame of edges.
#'
#' @examples
#' system.file("extdata", package = "pkgGraphR") |>
#'     buildPackageGraph()
#'
#' @export
buildPackageGraph <- function(x, unique.edges = TRUE, only.connected = FALSE){
    stopifnot(is.character(x), length(x) == 1)
    stopifnot(is.logical(unique.edges), length(unique.edges) == 1)
    funs <- collectFunNames(x)
    allFuns <- purrr::map(funs, \(x) x$name) |> unlist(use.names = FALSE) |> unique()
    res <- lapply(names(funs), \(z){
        findFunCalls(z, funs)
    }) |> purrr::compact() |> unlist()
    edges <- data.frame(from = unname(res), to = names(res)) |>
        dplyr::filter(!is.na(from), !is.na(to))
    res <- list(nodes = allFuns, edges = edges)
    if(unique.edges){
        res$edges <- dplyr::distinct(res$edges)
    }
    if(only.connected){
        drop <- setdiff(res$nodes, unique(c(res$edges$from, res$edges$to)))
        res$nodes <- setdiff(res$nodes, drop)
    }
    return(res)
}

#' Plot a graph or diagram of a package
#'
#' @description From a list of nodes and edges, plots a diagram or graph
#' @name plotPackageGraph
#'
#' @param graph A list generated by \code{\link[pkgGraphR]{buildPackageGraph}}
#' containing edges and nodes of the graph.
#' @param fun.list An optional list generated by \code{\link[pkgGraphR]{collectFunNames}}
#' containing each files function assignments. Used only if `use.subgraphs` or
#' `use.colors` are true
#' @param use.subgraphs Logical indicating whether the graph should be partitioned
#' into subgraphs by the file in which the function assignment was made.
#' DEFAULT: FALSE
#' @param use.colors Logical indicating whether the nodes of the graph should be
#' colored by the file in which the function assignment was made.
#' N.B. No legend is plotted for the colors.
#' DEFAULT: FALSE
#'
#' @seealso \code{\link[pkgGraphR]{collectFunNames}}, \code{\link[pkgGraphR]{buildPackageGraph}}
#'
#' @return A grviz plot.
#'
#' @examples
#' pkgGraph <- system.file("extdata", package = "pkgGraphR") |>
#'     buildPackageGraph()
#' plotPackageGraph(graph = pkgGraph)
#'
#' pkgFuns <- system.file("extdata", package = "pkgGraphR") |>
#'     collectFunNames()
#'
#' plotPackageGraph(graph = pkgGraph, fun.list = pkgFuns, use.subgraphs = TRUE)
#' plotPackageGraph(graph = pkgGraph, fun.list = pkgFuns, use.colors = TRUE)
#' plotPackageGraph(graph = pkgGraph, fun.list = pkgFuns, use.colors = TRUE, use.subgraphs = TRUE)
#'
#' @export
plotPackageGraph <- function(graph,
                             fun.list,
                             use.subgraphs = FALSE,
                             use.colors = FALSE){
    stopifnot(is.list(graph), all(c("nodes", "edges") %in% names(graph)))
    if((isTRUE(use.subgraphs) | isTRUE(use.colors)) & missing(fun.list)){
        stop("'fun.list' is needed to use colors or subgraphs")
    }
    edges <- sapply(1:nrow(graph$edges), \(i){
        paste0(graph$edges$from[i], " -> ", graph$edges$to[i])
    }) |> paste0(collapse = "\n") |>
        # "." is not allowed in grViz node names
        gsub(pattern = "\\.", replacement = "")
    nodes <- paste0(graph$nodes, collapse = ";") |> gsub(pattern = "\\.", replacement = "")

    graph.init <- paste0(
        "digraph{", "\n",
        "graph[rankdir = LR]", "\n",
        "node[shape = box, style = rounded, fontname = Helvetica]", "\n"
    )

    if(!missing(fun.list)){
        fun.list <- fun.list |>
            purrr::map(\(x) x[x$name %in% graph$nodes,]) |>
            purrr::keep(\(x) nrow(x) > 0)
    }

    if(isTRUE(use.subgraphs) & !isTRUE(use.colors)){
        subGraphNames <- names(fun.list) |>
            strsplit("/") |>
            purrr::map(~ purrr::pluck(.x, length(.x))) |>
            unlist()

        subGraphs <- lapply(1:length(subGraphNames), \(i){
            paste0(
                "subgraph cluster_", i, " {", " \n",
                "graph[shape = rectangle, style = rounded]", "\n",
                "label = '", subGraphNames[i], "' \n",
                paste0(gsub("\\.", "", fun.list[[i]]$name), collapse = "\n"),
                "\n}", "\n"
            )
        })
        fullGraph <- paste0(
            graph.init,
            paste0(unlist(subGraphs), collapse = " "),
            edges, "\n",
            "}"
        )
    } else if(isTRUE(use.colors) & !isTRUE(use.subgraphs)){
        graph.init <- paste0(
            "digraph{", "\n",
            "graph[rankdir = LR]", "\n"
        )
        nodeColors <- names(fun.list) |>
            strsplit("/") |>
            purrr::map(~ purrr::pluck(.x, length(.x))) |>
            unlist()
        nodeColors <- stats::setNames(colors51[1:length(nodeColors)], nodeColors)

        colorNodes <- lapply(1:length(nodeColors), \(i){
            paste0(
                "node [shape = box, style = \"rounded,filled\", fontname = Helvetica, fillcolor = '",
                nodeColors[i], "']\n",
                paste0(gsub("\\.", "", fun.list[[i]]$name),  collapse = "; ")
            )
        })
        fullGraph <- paste0(
            graph.init,
            paste0(unlist(colorNodes), collapse = "\n"), "\n",
            edges, "\n",
            "}"
        )
    } else if(isTRUE(use.colors) & isTRUE(use.subgraphs)){
        graph.init <- paste0(
            "digraph{", "\n",
            "graph[rankdir = LR]", "\n"
        )
        subGraphNames <- names(fun.list) |>
            strsplit("/") |>
            purrr::map(~ purrr::pluck(.x, length(.x))) |>
            unlist()
        nodeColors <- stats::setNames(colors51[1:length(subGraphNames)], subGraphNames)

        colorSubGraphs <- lapply(1:length(subGraphNames), \(i){
            paste0(
                "subgraph cluster_", i, " {", " \n",
                "graph[shape = rectangle, style = rounded]", "\n",
                "label = '", subGraphNames[i], "' \n",
                "node [shape = box, style = \"rounded,filled\", fontname = Helvetica, fillcolor = '",
                nodeColors[i], "']\n",
                paste0(gsub("\\.", "", fun.list[[i]]$name), collapse = "\n"),
                "\n}", "\n"
            )
        })
        fullGraph <- paste0(
            graph.init,
            paste0(unlist(colorSubGraphs), collapse = " "),
            edges, "\n",
            "}"
        )
    } else {
        fullGraph <- paste0(
            graph.init,
            nodes,
            edges, "\n",
            "}"
        )
    }
    p <- DiagrammeR::grViz(
        fullGraph
    )
    return(p)
}