images-tree.go

package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"

	"github.com/dotcloud/docker/api/client"
	"github.com/dotcloud/docker/engine"
	"github.com/dotcloud/docker/utils"
)

var flViz = flag.Bool("viz", false, "display digraph viz")

func main() {
	flag.Parse()

	cli := client.NewDockerCli(os.Stdin, os.Stdout, os.Stderr,
		os.Getenv("_DOCKER_CLIENT_PROTO"),
		os.Getenv("_DOCKER_CLIENT_ADDR"),
		nil)

	body, _, err := readBody(cli.Call("GET", "/images/json?all=1", nil, false))
	if err != nil {
		log.Fatal(err)
	}

	outs := engine.NewTable("Created", 0)
	if _, err := outs.ReadListFrom(body); err != nil {
		log.Fatal(err)
	}

	var (
		printNode  func(out io.Writer, noTrunc bool, image *engine.Env, prefix string)
		startImage *engine.Env

		roots    = engine.NewTable("Created", outs.Len())
		byParent = make(map[string]*engine.Table)
	)

	for _, image := range outs.Data {
		if image.Get("ParentId") == "" {
			roots.Add(image)
		} else {
			if children, exists := byParent[image.Get("ParentId")]; exists {
				children.Add(image)
			} else {
				byParent[image.Get("ParentId")] = engine.NewTable("Created", 1)
				byParent[image.Get("ParentId")].Add(image)
			}
		}
	}

	if *flViz {
		fmt.Printf("digraph docker {\n")
		printNode = printVizNode
	} else {
		printNode = printTreeNode
	}

	if startImage != nil {
		root := engine.NewTable("Created", 1)
		root.Add(startImage)
		WalkTree(os.Stdout, false, root, byParent, "", printNode)
	}
	if *flViz {
		fmt.Printf(" base [style=invisible]\n}\n")
	}
}

func WalkTree(out io.Writer, noTrunc bool, images *engine.Table, byParent map[string]*engine.Table, prefix string, printNode func(out io.Writer, noTrunc bool, image *engine.Env, prefix string)) {
	length := images.Len()
	if length > 1 {
		for index, image := range images.Data {
			if index+1 == length {
				printNode(out, noTrunc, image, prefix+"└─")
				if subimages, exists := byParent[image.Get("Id")]; exists {
					WalkTree(out, noTrunc, subimages, byParent, prefix+"  ", printNode)
				}
			} else {
				printNode(out, noTrunc, image, prefix+"\u251C─")
				if subimages, exists := byParent[image.Get("Id")]; exists {
					WalkTree(out, noTrunc, subimages, byParent, prefix+"\u2502 ", printNode)
				}
			}
		}
	} else {
		for _, image := range images.Data {
			printNode(out, noTrunc, image, prefix+"└─")
			if subimages, exists := byParent[image.Get("Id")]; exists {
				WalkTree(out, noTrunc, subimages, byParent, prefix+"  ", printNode)
			}
		}
	}
}

func printVizNode(out io.Writer, noTrunc bool, image *engine.Env, prefix string) {
	var (
		imageID  string
		parentID string
	)
	if noTrunc {
		imageID = image.Get("Id")
		parentID = image.Get("ParentId")
	} else {
		imageID = utils.TruncateID(image.Get("Id"))
		parentID = utils.TruncateID(image.Get("ParentId"))
	}
	if parentID == "" {
		fmt.Fprintf(out, " base -> \"%s\" [style=invis]\n", imageID)
	} else {
		fmt.Fprintf(out, " \"%s\" -> \"%s\"\n", parentID, imageID)
	}
	if image.GetList("RepoTags")[0] != "<none>:<none>" {
		fmt.Fprintf(out, " \"%s\" [label=\"%s\\n%s\",shape=box,fillcolor=\"paleturquoise\",style=\"filled,rounded\"];\n",
			imageID, imageID, strings.Join(image.GetList("RepoTags"), "\\n"))
	}
}

// FIXME: --viz and --tree are deprecated. Remove them in a future version.
func printTreeNode(out io.Writer, noTrunc bool, image *engine.Env, prefix string) {
	var imageID string
	if noTrunc {
		imageID = image.Get("Id")
	} else {
		imageID = utils.TruncateID(image.Get("Id"))
	}

	fmt.Fprintf(out, "%s%s Virtual Size: %s", prefix, imageID, utils.HumanSize(image.GetInt64("VirtualSize")))
	if image.GetList("RepoTags")[0] != "<none>:<none>" {
		fmt.Fprintf(out, " Tags: %s\n", strings.Join(image.GetList("RepoTags"), ", "))
	} else {
		fmt.Fprint(out, "\n")
	}
}

func readBody(stream io.ReadCloser, statusCode int, err error) ([]byte, int, error) {
	if stream != nil {
		defer stream.Close()
	}
	if err != nil {
		return nil, statusCode, err
	}
	body, err := ioutil.ReadAll(stream)
	if err != nil {
		return nil, -1, err
	}
	return body, statusCode, nil
}