<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	// "github.com/davecgh/go-spew/spew"
	"github.com/vharitonsky/iniflags"
	"github.com/zubairhamed/canopus"
	"os"
	"time"
)

// types to unmarshal json data from tradfri
// partially generated with https://mholt.github.io/json-to-go/
// struct names derived from https://github.com/IPSO-Alliance/pub/blob/master/reg/README.md
type device_ids []int
type group_ids []int

type device_desc struct {
	Device struct {
		Manufacturer          string `json:"0"`
		DeviceDescription     string `json:"1"`
		SerialNumber          string `json:"2"`
		FirmwareVersion       string `json:"3"`
		AvailablePowerSources int    `json:"6"`
	} `json:"3"`
	LightControl []struct {
		Power   int `json:"5850"`
		Dim     int `json:"5851"`
		Num9003 int `json:"9003"`
	} `json:"3311"`
	ApplicationType int    `json:"5750"`
	DeviceName      string `json:"9001"`
	Num9002         int    `json:"9002"`
	DeviceID        int    `json:"9003"`
	Num9019         int    `json:"9019"`
	Num9020         int    `json:"9020"`
	Num9054         int    `json:"9054"`
}

type group_desc struct {
	Power     int    `json:"5850"`
	Dim       int    `json:"5851"`
	GroupName string `json:"9001"`
	Num9002   int    `json:"9002"`
	GroupID   int    `json:"9003"`
	Num9018   struct {
		Num15002 struct {
			Num9003 []int `json:"9003"`
		} `json:"15002"`
	} `json:"9018"`
	Num9039 int `json:"9039"`
}

// type to read the config file

type tradfri_cfg struct {
	Hubip string
	Key   string
}

// flags
var (
	gateway     = flag.String("gateway", "127.0.0.1", "Address of Tradfri gateway.")
	key         = flag.String("key", "deadbeef", "API key to access gateway.")
	action      = flag.String("action", "status", "action to take [dim|status|power]).")
	target      = flag.Int("target", 0, "Target value (0-100 for dim, 0 or 1 for power).")
	target_id   = flag.Int("id", 65537, "Device or Group ID.")
	target_name = flag.String("name", "noname", "Device or Group name")
	steps       = flag.Int("steps", 10, "Number of intermediate steps for dim action.")
	period      = flag.Int("period", 60, "Time period in seconds to run dim action over.")
)

// process flags
func init() {
	flag.Usage = usage
	iniflags.SetConfigFile("tradfri.ini")
	iniflags.Parse()
}

// usage info
func usage() {
	flag.PrintDefaults()
	os.Exit(1)
}

// deal with errors gracelessly
func check(e error) {
	if e != nil {
		panic(e)
	}
}

func tradfri_conn(address string, key string) canopus.Connection {
	var tradfri_gw = fmt.Sprintf("%s:5684", address)
	fmt.Println("Connecting to tradfri gateway... ")
	conn, err := canopus.DialDTLS(tradfri_gw, "", key)
	check(err)
	fmt.Println("connected")
	return conn
}

func list_device_ids(conn canopus.Connection) device_ids {
	var device_id_list device_ids

	// setup request for device ids
	req := canopus.NewRequest(canopus.MessageConfirmable, canopus.Get)
	req.SetStringPayload("")
	req.SetRequestURI("/15001")

	// request device ids
	fmt.Print("Looking for devices... ")
	resp, err := conn.Send(req)
	check(err)

	json.Unmarshal([]byte(resp.GetMessage().GetPayload().String()), &device_id_list)
	return device_id_list
}

func list_group_ids(conn canopus.Connection) group_ids {
	var group_id_list group_ids

	// setup request for device ids
	req := canopus.NewRequest(canopus.MessageConfirmable, canopus.Get)
	req.SetStringPayload("")
	req.SetRequestURI("/15004")

	// request device ids
	fmt.Print("Looking for groups... ")
	resp, err := conn.Send(req)
	check(err)

	json.Unmarshal([]byte(resp.GetMessage().GetPayload().String()), &group_id_list)
	return group_id_list
}

func get_group_info(group_id int, conn canopus.Connection) {
	req := canopus.NewRequest(canopus.MessageConfirmable, canopus.Get)
	req.SetStringPayload("")
	ru := fmt.Sprintf("/15004/%v", group_id)
	req.SetRequestURI(ru)
	dresp, err := conn.Send(req)
	check(err)

	// output basic device information
	var desc group_desc
	json.Unmarshal([]byte(dresp.GetMessage().GetPayload().String()), &desc)
	fmt.Printf("ID: %v, Name: %v\n", desc.GroupID, desc.GroupName)
	fmt.Printf("Power: %v, Dim: %v\n", desc.Power, desc.Dim)

}

func list_groups(group_id_list group_ids, conn canopus.Connection) {
	// enumerating group information
	fmt.Println("enumerating:")
	for _, group := range group_id_list {
		get_group_info(group, conn)
		// sleep for a while to avoid flood protection
		time.Sleep(500 * time.Millisecond)
	}
}

func get_device_info(device_id int, conn canopus.Connection) {
	req := canopus.NewRequest(canopus.MessageConfirmable, canopus.Get)
	req.SetStringPayload("")
	ru := fmt.Sprintf("/15001/%v", device_id)
	req.SetRequestURI(ru)
	dresp, err := conn.Send(req)
	check(err)

	// output basic device information
	var desc device_desc
	json.Unmarshal([]byte(dresp.GetMessage().GetPayload().String()), &desc)
	fmt.Printf("ID: %v, Name; %v, Description: %v\n",
		desc.DeviceID, desc.DeviceName, desc.Device.DeviceDescription)

	// only output light control info if available
	if len(desc.LightControl) > 0 {
		for count, entry := range desc.LightControl {
			fmt.Printf("Light Control Set %v, Power: %v, Dim: %v\n",
				count, entry.Power, entry.Dim)
		}
	} else {
		fmt.Println("No light control values")
	}
}

func list_devices(device_id_list device_ids, conn canopus.Connection) {
	fmt.Println("enumerating:")
	for _, device := range device_id_list {
		get_device_info(device, conn)

		// sleep for a while to avoid flood protection
		time.Sleep(500 * time.Millisecond)
	}
}

func power_device(device_id int, val int, conn canopus.Connection) {
	get_device_info(device_id, conn)
	req := canopus.NewRequest(canopus.MessageConfirmable, canopus.Put)
	payload := fmt.Sprintf("{ \"3311\" : [{ \"5850\" : %v }] }", val)
	req.SetStringPayload(payload)
	ru := fmt.Sprintf("/15001/%v", device_id)
	req.SetRequestURI(ru)
	_, err := conn.Send(req)
	check(err)
	get_device_info(device_id, conn)
}

func dim_device(device_id int, val int, conn canopus.Connection) {
	get_device_info(device_id, conn)
	req := canopus.NewRequest(canopus.MessageConfirmable, canopus.Put)
	payload := fmt.Sprintf("{ \"3311\" : [{ \"5851\" : %v }] }", val)
	req.SetStringPayload(payload)
	ru := fmt.Sprintf("/15001/%v", device_id)
	req.SetRequestURI(ru)
	_, err := conn.Send(req)
	check(err)
	get_device_info(device_id, conn)
}

func power_group(group_id int, val int, conn canopus.Connection) {
	get_group_info(group_id, conn)
	req := canopus.NewRequest(canopus.MessageConfirmable, canopus.Put)
	payload := fmt.Sprintf("{ \"5850\": %d }", val)
	req.SetStringPayload(payload)
	ru := fmt.Sprintf("/15004/%v", group_id)
	req.SetRequestURI(ru)
	_, err := conn.Send(req)
	check(err)
	get_group_info(group_id, conn)
}

func dim_group(group_id int, val int, conn canopus.Connection) {
	get_group_info(group_id, conn)
	req := canopus.NewRequest(canopus.MessageConfirmable, canopus.Put)
	payload := fmt.Sprintf("{ \"5851\": %d }", val)
	req.SetStringPayload(payload)
	ru := fmt.Sprintf("/15004/%v", group_id)
	req.SetRequestURI(ru)
	_, err := conn.Send(req)
	check(err)
	get_group_info(group_id, conn)
}

func main() {
	conn := tradfri_conn(*gateway, *key)
	if *action == "status" {
		list_devices(list_device_ids(conn), conn)
		list_groups(list_group_ids(conn), conn)
	}
	if *action == "power" {
		fmt.Printf("power target %v on id %v", *target, *target_id)
	}
	if *action == "dim" {
		fmt.Printf("dim target %v on id %v", *target, *target_id)
	}
}
<commit_msg>twiddling with flags<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	// "github.com/davecgh/go-spew/spew"
	"github.com/vharitonsky/iniflags"
	"github.com/zubairhamed/canopus"
	"os"
	"time"
)

// types to unmarshal json data from tradfri
// partially generated with https://mholt.github.io/json-to-go/
// struct names derived from https://github.com/IPSO-Alliance/pub/blob/master/reg/README.md
type device_ids []int
type group_ids []int

type device_desc struct {
	Device struct {
		Manufacturer          string `json:"0"`
		DeviceDescription     string `json:"1"`
		SerialNumber          string `json:"2"`
		FirmwareVersion       string `json:"3"`
		AvailablePowerSources int    `json:"6"`
	} `json:"3"`
	LightControl []struct {
		Power   int `json:"5850"`
		Dim     int `json:"5851"`
		Num9003 int `json:"9003"`
	} `json:"3311"`
	ApplicationType int    `json:"5750"`
	DeviceName      string `json:"9001"`
	Num9002         int    `json:"9002"`
	DeviceID        int    `json:"9003"`
	Num9019         int    `json:"9019"`
	Num9020         int    `json:"9020"`
	Num9054         int    `json:"9054"`
}

type group_desc struct {
	Power     int    `json:"5850"`
	Dim       int    `json:"5851"`
	GroupName string `json:"9001"`
	Num9002   int    `json:"9002"`
	GroupID   int    `json:"9003"`
	Num9018   struct {
		Num15002 struct {
			Num9003 []int `json:"9003"`
		} `json:"15002"`
	} `json:"9018"`
	Num9039 int `json:"9039"`
}

// type to read the config file

type tradfri_cfg struct {
	Hubip string
	Key   string
}

// flags
var (
	gateway     = flag.String("gateway", "127.0.0.1", "Address of Tradfri gateway.")
	key         = flag.String("key", "deadbeef", "API key to access gateway.")
	action      = flag.String("action", "status", "action to take [dim|status|power]).")
	target      = flag.Int("target", 0, "Target value (0-100 for dim, 0 or 1 for power).")
	target_id   = flag.Int("id", 65537, "Device or Group ID.")
	target_name = flag.String("name", "", "Device or Group name")
	steps       = flag.Int("steps", 10, "Number of intermediate steps for dim action.")
	period      = flag.Int("period", 60, "Time period in seconds to run dim action over.")
)

// process flags
func init() {
	flag.Usage = usage
	iniflags.SetConfigFile("tradfri.ini")
	iniflags.Parse()
}

// usage info
func usage() {
	flag.PrintDefaults()
	os.Exit(1)
}

// deal with errors gracelessly
func check(e error) {
	if e != nil {
		panic(e)
	}
}

func tradfri_conn(address string, key string) canopus.Connection {
	var tradfri_gw = fmt.Sprintf("%s:5684", address)
	fmt.Println("Connecting to tradfri gateway... ")
	conn, err := canopus.DialDTLS(tradfri_gw, "", key)
	check(err)
	fmt.Println("connected")
	return conn
}

func list_device_ids(conn canopus.Connection) device_ids {
	var device_id_list device_ids

	// setup request for device ids
	req := canopus.NewRequest(canopus.MessageConfirmable, canopus.Get)
	req.SetStringPayload("")
	req.SetRequestURI("/15001")

	// request device ids
	fmt.Print("Looking for devices... ")
	resp, err := conn.Send(req)
	check(err)

	json.Unmarshal([]byte(resp.GetMessage().GetPayload().String()), &device_id_list)
	return device_id_list
}

func list_group_ids(conn canopus.Connection) group_ids {
	var group_id_list group_ids

	// setup request for device ids
	req := canopus.NewRequest(canopus.MessageConfirmable, canopus.Get)
	req.SetStringPayload("")
	req.SetRequestURI("/15004")

	// request device ids
	fmt.Print("Looking for groups... ")
	resp, err := conn.Send(req)
	check(err)

	json.Unmarshal([]byte(resp.GetMessage().GetPayload().String()), &group_id_list)
	return group_id_list
}

func get_group_info(group_id int, conn canopus.Connection) {
	req := canopus.NewRequest(canopus.MessageConfirmable, canopus.Get)
	req.SetStringPayload("")
	ru := fmt.Sprintf("/15004/%v", group_id)
	req.SetRequestURI(ru)
	dresp, err := conn.Send(req)
	check(err)

	// output basic device information
	var desc group_desc
	json.Unmarshal([]byte(dresp.GetMessage().GetPayload().String()), &desc)
	fmt.Printf("ID: %v, Name: %v\n", desc.GroupID, desc.GroupName)
	fmt.Printf("Power: %v, Dim: %v\n", desc.Power, desc.Dim)

}

func list_groups(group_id_list group_ids, conn canopus.Connection) {
	// enumerating group information
	fmt.Println("enumerating:")
	for _, group := range group_id_list {
		get_group_info(group, conn)
		// sleep for a while to avoid flood protection
		time.Sleep(500 * time.Millisecond)
	}
}

func get_device_info(device_id int, conn canopus.Connection) {
	req := canopus.NewRequest(canopus.MessageConfirmable, canopus.Get)
	req.SetStringPayload("")
	ru := fmt.Sprintf("/15001/%v", device_id)
	req.SetRequestURI(ru)
	dresp, err := conn.Send(req)
	check(err)

	// output basic device information
	var desc device_desc
	json.Unmarshal([]byte(dresp.GetMessage().GetPayload().String()), &desc)
	fmt.Printf("ID: %v, Name; %v, Description: %v\n",
		desc.DeviceID, desc.DeviceName, desc.Device.DeviceDescription)

	// only output light control info if available
	if len(desc.LightControl) > 0 {
		for count, entry := range desc.LightControl {
			fmt.Printf("Light Control Set %v, Power: %v, Dim: %v\n",
				count, entry.Power, entry.Dim)
		}
	} else {
		fmt.Println("No light control values")
	}
}

func list_devices(device_id_list device_ids, conn canopus.Connection) {
	fmt.Println("enumerating:")
	for _, device := range device_id_list {
		get_device_info(device, conn)

		// sleep for a while to avoid flood protection
		time.Sleep(500 * time.Millisecond)
	}
}

func power_device(device_id int, val int, conn canopus.Connection) {
	get_device_info(device_id, conn)
	req := canopus.NewRequest(canopus.MessageConfirmable, canopus.Put)
	payload := fmt.Sprintf("{ \"3311\" : [{ \"5850\" : %v }] }", val)
	req.SetStringPayload(payload)
	ru := fmt.Sprintf("/15001/%v", device_id)
	req.SetRequestURI(ru)
	_, err := conn.Send(req)
	check(err)
	get_device_info(device_id, conn)
}

func dim_device(device_id int, val int, conn canopus.Connection) {
	get_device_info(device_id, conn)
	req := canopus.NewRequest(canopus.MessageConfirmable, canopus.Put)
	payload := fmt.Sprintf("{ \"3311\" : [{ \"5851\" : %v }] }", val)
	req.SetStringPayload(payload)
	ru := fmt.Sprintf("/15001/%v", device_id)
	req.SetRequestURI(ru)
	_, err := conn.Send(req)
	check(err)
	get_device_info(device_id, conn)
}

func power_group(group_id int, val int, conn canopus.Connection) {
	get_group_info(group_id, conn)
	req := canopus.NewRequest(canopus.MessageConfirmable, canopus.Put)
	payload := fmt.Sprintf("{ \"5850\": %d }", val)
	req.SetStringPayload(payload)
	ru := fmt.Sprintf("/15004/%v", group_id)
	req.SetRequestURI(ru)
	_, err := conn.Send(req)
	check(err)
	get_group_info(group_id, conn)
}

func dim_group(group_id int, val int, conn canopus.Connection) {
	get_group_info(group_id, conn)
	req := canopus.NewRequest(canopus.MessageConfirmable, canopus.Put)
	payload := fmt.Sprintf("{ \"5851\": %d }", val)
	req.SetStringPayload(payload)
	ru := fmt.Sprintf("/15004/%v", group_id)
	req.SetRequestURI(ru)
	_, err := conn.Send(req)
	check(err)
	get_group_info(group_id, conn)
}

func main() {
	conn := tradfri_conn(*gateway, *key)
	if *action == "status" {
		list_devices(list_device_ids(conn), conn)
		list_groups(list_group_ids(conn), conn)
	}
	if *action == "power" {
		fmt.Printf("power target %v on id %v", *target, *target_id)
	}
	if *action == "dim" {
		fmt.Printf("dim target %v on id %v", *target, *target_id)
	}
}
<|endoftext|>


<commit_before>// Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

// Package gi18n implements internationalization and localization.
package gi18n

var (
	// defaultManager is the default i18n instance for package functions.
	defaultManager = Instance()
)

// SetPath sets the directory path storing i18n files.
func SetPath(path string) error {
	return defaultManager.SetPath(path)
}

// SetLanguage sets the language for translator.
func SetLanguage(language string) {
	defaultManager.SetLanguage(language)
}

// SetDelimiters sets the delimiters for translator.
func SetDelimiters(left, right string) {
	defaultManager.SetDelimiters(left, right)
}

// T is alias of Translate for convenience.
func T(content string, language ...string) string {
	return defaultManager.T(content, language...)
}

// Tf is alias of TranslateFormat for convenience.
func Tf(format string, values ...interface{}) string {
	return defaultManager.TranslateFormat(format, values...)
}

// Tfl is alias of TranslateFormatLang for convenience.
func Tfl(language string, format string, values ...interface{}) string {
	return defaultManager.TranslateFormatLang(language, format, values...)
}

// TranslateFormat translates, formats and returns the <format> with configured language
// and given <values>.
func TranslateFormat(format string, values ...interface{}) string {
	return defaultManager.TranslateFormat(format, values...)
}

// TranslateFormatLang translates, formats and returns the <format> with configured language
// and given <values>. The parameter <language> specifies custom translation language ignoring
// configured language. If <language> is given empty string, it uses the default configured
// language for the translation.
func TranslateFormatLang(language string, format string, values ...interface{}) string {
	return defaultManager.TranslateFormatLang(language, format, values...)
}

// Translate translates <content> with configured language and returns the translated content.
// The parameter <language> specifies custom translation language ignoring configured language.
func Translate(content string, language ...string) string {
	return defaultManager.Translate(content, language...)
}

// GetValue retrieves and returns the configured content for given key and specified language.
// It returns an empty string if not found.
func GetContent(key string, language ...string) string {
	return defaultManager.GetContent(key, language...)
}
<commit_msg>improve package gi18n<commit_after>// Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

// Package gi18n implements internationalization and localization.
package gi18n

// SetPath sets the directory path storing i18n files.
func SetPath(path string) error {
	return Instance().SetPath(path)
}

// SetLanguage sets the language for translator.
func SetLanguage(language string) {
	Instance().SetLanguage(language)
}

// SetDelimiters sets the delimiters for translator.
func SetDelimiters(left, right string) {
	Instance().SetDelimiters(left, right)
}

// T is alias of Translate for convenience.
func T(content string, language ...string) string {
	return Instance().T(content, language...)
}

// Tf is alias of TranslateFormat for convenience.
func Tf(format string, values ...interface{}) string {
	return Instance().TranslateFormat(format, values...)
}

// Tfl is alias of TranslateFormatLang for convenience.
func Tfl(language string, format string, values ...interface{}) string {
	return Instance().TranslateFormatLang(language, format, values...)
}

// TranslateFormat translates, formats and returns the <format> with configured language
// and given <values>.
func TranslateFormat(format string, values ...interface{}) string {
	return Instance().TranslateFormat(format, values...)
}

// TranslateFormatLang translates, formats and returns the <format> with configured language
// and given <values>. The parameter <language> specifies custom translation language ignoring
// configured language. If <language> is given empty string, it uses the default configured
// language for the translation.
func TranslateFormatLang(language string, format string, values ...interface{}) string {
	return Instance().TranslateFormatLang(language, format, values...)
}

// Translate translates <content> with configured language and returns the translated content.
// The parameter <language> specifies custom translation language ignoring configured language.
func Translate(content string, language ...string) string {
	return Instance().Translate(content, language...)
}

// GetValue retrieves and returns the configured content for given key and specified language.
// It returns an empty string if not found.
func GetContent(key string, language ...string) string {
	return Instance().GetContent(key, language...)
}
<|endoftext|>


<commit_before>// Package trigram is a dumb trigram index
package trigram

// T is a trigram
type T uint32

func (t T) String() string {
	b := [3]byte{byte(t >> 16), byte(t >> 8), byte(t)}
	return string(b[:])
}

// DocID is a document ID
type DocID int

// Index is a trigram index
type Index map[T][]DocID

// Extract returns a list of trigrams in s
func Extract(s string, trigrams []T) []T {

	for i := 0; i <= len(s)-3; i++ {
		t := T(uint32(s[i])<<16 | uint32(s[i+1])<<8 | uint32(s[i+2]))
		trigrams = appendIfUnique(trigrams, t)
	}

	return trigrams
}

func appendIfUnique(t []T, n T) []T {
	for _, v := range t {
		if v == n {
			return t
		}
	}

	return append(t, n)
}

// NewIndex returns an index for the strings in docs
func NewIndex(docs []string) Index {

	idx := make(Index)

	var trigrams []T

	for id, d := range docs {
		ts := Extract(d, trigrams)
		docid := DocID(id)
		for _, t := range ts {
			idx[t] = append(idx[t], docid)
		}
		trigrams = trigrams[:0]
	}

	return idx
}

// Add adds a new string to the search index
func (idx Index) Add(s string) {

	id := DocID(len(idx))

	ts := Extract(s, nil)
	for _, t := range ts {
		idx[t] = append(idx[t], id)
	}
}

// Query returns a list of document IDs that match the trigrams in the query s
func (idx Index) Query(s string) []DocID {
	ts := Extract(s, nil)
	return idx.QueryTrigrams(ts)
}

// QueryTrigrams returns a list of document IDs that match the trigram set ts
func (idx Index) QueryTrigrams(ts []T) []DocID {

	midx := 0
	mtri := ts[midx]

	for i, t := range ts {
		if len(idx[t]) < len(idx[mtri]) {
			midx = i
			mtri = t
		}
	}

	ts[0], ts[midx] = ts[midx], ts[0]

	return idx.Filter(idx[mtri], ts[1:]...)
}

// Filter removes documents that don't contain the specified trigrams
func (idx Index) Filter(docs []DocID, ts ...T) []DocID {
	for _, t := range ts {
		docs = intersect(docs, idx[t])
	}

	return docs
}

func intersect(a, b []DocID) []DocID {

	// TODO(dgryski): reduce allocations by reusing A

	var aidx, bidx int

	var result []DocID

scan:
	for aidx < len(a) && bidx < len(b) {
		if a[aidx] == b[bidx] {
			result = append(result, a[aidx])
			aidx++
			bidx++
			if aidx >= len(a) || bidx >= len(b) {
				break scan
			}
		}

		for a[aidx] < b[bidx] {
			aidx++
			if aidx >= len(a) {
				break scan
			}
		}

		for bidx < len(b) && a[aidx] > b[bidx] {
			bidx++
			if bidx >= len(b) {
				break scan
			}
		}
	}

	return result
}
<commit_msg>Replace DocID with a uint32<commit_after>// Package trigram is a dumb trigram index
package trigram

// T is a trigram
type T uint32

func (t T) String() string {
	b := [3]byte{byte(t >> 16), byte(t >> 8), byte(t)}
	return string(b[:])
}

// DocID is a document ID
type DocID uint32

// Index is a trigram index
type Index map[T][]DocID

// Extract returns a list of trigrams in s
func Extract(s string, trigrams []T) []T {

	for i := 0; i <= len(s)-3; i++ {
		t := T(uint32(s[i])<<16 | uint32(s[i+1])<<8 | uint32(s[i+2]))
		trigrams = appendIfUnique(trigrams, t)
	}

	return trigrams
}

func appendIfUnique(t []T, n T) []T {
	for _, v := range t {
		if v == n {
			return t
		}
	}

	return append(t, n)
}

// NewIndex returns an index for the strings in docs
func NewIndex(docs []string) Index {

	idx := make(Index)

	var trigrams []T

	for id, d := range docs {
		ts := Extract(d, trigrams)
		docid := DocID(id)
		for _, t := range ts {
			idx[t] = append(idx[t], docid)
		}
		trigrams = trigrams[:0]
	}

	return idx
}

// Add adds a new string to the search index
func (idx Index) Add(s string) {

	id := DocID(len(idx))

	ts := Extract(s, nil)
	for _, t := range ts {
		idx[t] = append(idx[t], id)
	}
}

// Query returns a list of document IDs that match the trigrams in the query s
func (idx Index) Query(s string) []DocID {
	ts := Extract(s, nil)
	return idx.QueryTrigrams(ts)
}

// QueryTrigrams returns a list of document IDs that match the trigram set ts
func (idx Index) QueryTrigrams(ts []T) []DocID {

	midx := 0
	mtri := ts[midx]

	for i, t := range ts {
		if len(idx[t]) < len(idx[mtri]) {
			midx = i
			mtri = t
		}
	}

	ts[0], ts[midx] = ts[midx], ts[0]

	return idx.Filter(idx[mtri], ts[1:]...)
}

// Filter removes documents that don't contain the specified trigrams
func (idx Index) Filter(docs []DocID, ts ...T) []DocID {
	for _, t := range ts {
		docs = intersect(docs, idx[t])
	}

	return docs
}

func intersect(a, b []DocID) []DocID {

	// TODO(dgryski): reduce allocations by reusing A

	var aidx, bidx int

	var result []DocID

scan:
	for aidx < len(a) && bidx < len(b) {
		if a[aidx] == b[bidx] {
			result = append(result, a[aidx])
			aidx++
			bidx++
			if aidx >= len(a) || bidx >= len(b) {
				break scan
			}
		}

		for a[aidx] < b[bidx] {
			aidx++
			if aidx >= len(a) {
				break scan
			}
		}

		for bidx < len(b) && a[aidx] > b[bidx] {
			bidx++
			if bidx >= len(b) {
				break scan
			}
		}
	}

	return result
}
<|endoftext|>


<commit_before>/* Package trigram is a dumb trigram index */
package trigram

type tri uint32

type Index map[tri][]int

func Extract(s string, trigrams []tri) []tri {

	for i := 0; i <= len(s)-3; i++ {
		t := tri(uint32(s[i])<<16 | uint32(s[i+1])<<8 | uint32(s[i+2]))
		trigrams = appendIfUnique(trigrams, t)
	}

	return trigrams
}

func appendIfUnique(t []tri, n tri) []tri {
	for _, v := range t {
		if v == n {
			return t
		}
	}

	return append(t, n)
}

func NewIndex(docs []string) Index {

	idx := make(Index)

	var trigrams []tri

	for id, d := range docs {
		ts := Extract(d, trigrams)
		for _, t := range ts {
			idx[t] = append(idx[t], id)
		}
		trigrams = trigrams[:0]
	}

	return idx
}

func (idx Index) Add(s string) {

	id := len(idx)

	ts := Extract(s, nil)
	for _, t := range ts {
		idx[t] = append(idx[t], id)
	}
}

func (idx Index) Query(s string) []int {
	ts := Extract(s, nil)
	return idx.QueryTrigrams(ts)
}

func (idx Index) QueryTrigrams(ts []tri) []int {

	midx := 0
	mtri := ts[midx]

	for i, t := range ts {
		if len(idx[t]) < len(idx[mtri]) {
			midx = i
			mtri = t
		}
	}

	ts[0], ts[midx] = ts[midx], ts[0]

	return idx.Filter(idx[mtri], ts[1:]...)
}

func (idx Index) Filter(docs []int, ts ...tri) []int {
	for _, t := range ts {
		docs = intersect(docs, idx[t])
	}

	return docs
}

func intersect(a, b []int) []int {

	// TODO(dgryski): reduce allocations by reusing A

	var aidx, bidx int

	var result []int

	for aidx < len(a) && bidx < len(b) {
		switch {
		case a[aidx] == b[bidx]:
			result = append(result, a[aidx])
			aidx++
			bidx++
		case a[aidx] < b[bidx]:
			aidx++
		case a[aidx] > b[bidx]:
			bidx++
		}
	}

	return result
}
<commit_msg>Export trigram type<commit_after>/* Package trigram is a dumb trigram index */
package trigram

type T uint32

type Index map[T][]int

func Extract(s string, trigrams []T) []T {

	for i := 0; i <= len(s)-3; i++ {
		t := T(uint32(s[i])<<16 | uint32(s[i+1])<<8 | uint32(s[i+2]))
		trigrams = appendIfUnique(trigrams, t)
	}

	return trigrams
}

func appendIfUnique(t []T, n T) []T {
	for _, v := range t {
		if v == n {
			return t
		}
	}

	return append(t, n)
}

func NewIndex(docs []string) Index {

	idx := make(Index)

	var trigrams []T

	for id, d := range docs {
		ts := Extract(d, trigrams)
		for _, t := range ts {
			idx[t] = append(idx[t], id)
		}
		trigrams = trigrams[:0]
	}

	return idx
}

func (idx Index) Add(s string) {

	id := len(idx)

	ts := Extract(s, nil)
	for _, t := range ts {
		idx[t] = append(idx[t], id)
	}
}

func (idx Index) Query(s string) []int {
	ts := Extract(s, nil)
	return idx.QueryTrigrams(ts)
}

func (idx Index) QueryTrigrams(ts []T) []int {

	midx := 0
	mtri := ts[midx]

	for i, t := range ts {
		if len(idx[t]) < len(idx[mtri]) {
			midx = i
			mtri = t
		}
	}

	ts[0], ts[midx] = ts[midx], ts[0]

	return idx.Filter(idx[mtri], ts[1:]...)
}

func (idx Index) Filter(docs []int, ts ...T) []int {
	for _, t := range ts {
		docs = intersect(docs, idx[t])
	}

	return docs
}

func intersect(a, b []int) []int {

	// TODO(dgryski): reduce allocations by reusing A

	var aidx, bidx int

	var result []int

	for aidx < len(a) && bidx < len(b) {
		switch {
		case a[aidx] == b[bidx]:
			result = append(result, a[aidx])
			aidx++
			bidx++
		case a[aidx] < b[bidx]:
			aidx++
		case a[aidx] > b[bidx]:
			bidx++
		}
	}

	return result
}
<|endoftext|>


<commit_before>package main

//#include<stdlib.h>
import "C"

import (
	"fmt"
	"github.com/mortdeus/egles/egl"
	"github.com/mortdeus/egles/gles/2"
)

var (
	attr = [...]egl.Int{
		egl.WIDTH, 500,
		egl.HEIGHT, 500,
		egl.NONE,
	}
	numConfig egl.Int
	max       egl.Int
	min       egl.Int
	ctx       egl.Context
	pbuf      egl.Surface
	configs   egl.Config
	b         egl.Boolean
)

func run() {
	disp := egl.GetDisplay(egl.DEFAULT_DISPLAY)
	defer egl.Terminate(disp)

	if ok := egl.Initialize(disp, &max, &min); !ok {
		panic("Initialize() failed")
		return
	}

	fmt.Printf("EGL Version: %v, %v\n", max, min)

	if ok := egl.GetConfigs(disp, nil, 0, &numConfig); !ok {
		panic("GetConfigs() failed")
	}
	configs := make([]egl.Config, int(numConfig))

	if ok := egl.GetConfigs(disp, &configs[0], numConfig, &numConfig); !ok {
		panic("GetConfigs() failed")
	}

	egl.BindAPI(egl.OPENGL_ES_API)
	ctx = egl.CreateContext(disp, configs[0], egl.NO_CONTEXT, nil)
	if ctx == egl.NO_CONTEXT {
		panic("CreateContext() failed")
		return
	}

	pbuf = egl.CreatePbufferSurface(disp, configs[0], &attr[0])
	configs = nil
	if ok := egl.MakeCurrent(disp, pbuf, pbuf, ctx); !ok {
		panic("MakeCurrent() failed")
		return
	}
	if ok := egl.MakeCurrent(disp, egl.NO_SURFACE, egl.NO_SURFACE, ctx); !ok {
		panic("MakeCurrent() failed")
		return
	}
	_ = egl.DestroySurface(disp, pbuf)
	_ = egl.DestroyContext(disp, ctx)

	return
}

func main() {
	run()
	println("Done")
}
<commit_msg>removed examples/test.go<commit_after><|endoftext|>


<commit_before>package fakes

import (
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	boshsys "github.com/cloudfoundry/bosh-utils/system"
)

type FakeCmdRunner struct {
	commandResults     map[string][]FakeCmdResult
	commandResultsLock sync.Mutex

	processes     map[string][]*FakeProcess
	processesLock sync.Mutex

	RunComplexCommands   []boshsys.Command
	RunCommands          [][]string
	RunCommandsWithInput [][]string
	RunCommandsQuietly   [][]string
	runCommandCallbacks  map[string]FakeCmdCallback

	CommandExistsValue bool
	AvailableCommands  map[string]bool
}

type FakeCmdCallback func()

type FakeCmdResult struct {
	Stdout     string
	Stderr     string
	ExitStatus int
	Error      error

	Sticky bool // Set to true if this result should ALWAYS be returned for the given command
}

type FakeProcess struct {
	StartErr error

	WaitCh chan boshsys.Result

	Waited     bool
	WaitResult boshsys.Result

	TerminatedNicely               bool
	TerminatedNicelyCallBack       func(*FakeProcess)
	TerminateNicelyKillGracePeriod time.Duration
	TerminateNicelyErr             error

	Stdout io.Writer
	Stderr io.Writer
}

func (p *FakeProcess) Wait() <-chan boshsys.Result {
	if p.Waited {
		panic("Cannot Wait() on process multiple times")
	}

	p.Waited = true
	p.WaitCh = make(chan boshsys.Result, 1)

	if p.TerminatedNicelyCallBack == nil {
		p.WaitCh <- p.WaitResult
	}
	return p.WaitCh
}

func (p *FakeProcess) TerminateNicely(killGracePeriod time.Duration) error {
	p.TerminateNicelyKillGracePeriod = killGracePeriod
	p.TerminatedNicely = true
	if p.TerminatedNicelyCallBack != nil {
		p.TerminatedNicelyCallBack(p)
	}
	return p.TerminateNicelyErr
}

func NewFakeCmdRunner() *FakeCmdRunner {
	return &FakeCmdRunner{
		AvailableCommands:   map[string]bool{},
		commandResults:      map[string][]FakeCmdResult{},
		runCommandCallbacks: map[string]FakeCmdCallback{},
		processes:           map[string][]*FakeProcess{},
	}
}

func (r *FakeCmdRunner) RunComplexCommand(cmd boshsys.Command) (string, string, int, error) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	r.RunComplexCommands = append(r.RunComplexCommands, cmd)

	runCmd := append([]string{cmd.Name}, cmd.Args...)

	r.runCallbackForCmd(runCmd)

	stdout, stderr, exitstatus, err := r.getOutputsForCmd(runCmd)

	if cmd.Stdout != nil {
		cmd.Stdout.Write([]byte(stdout))
	}

	if cmd.Stderr != nil {
		cmd.Stderr.Write([]byte(stderr))
	}

	return stdout, stderr, exitstatus, err
}

func (r *FakeCmdRunner) RunComplexCommandAsync(cmd boshsys.Command) (boshsys.Process, error) {
	r.processesLock.Lock()
	defer r.processesLock.Unlock()

	r.RunComplexCommands = append(r.RunComplexCommands, cmd)

	runCmd := append([]string{cmd.Name}, cmd.Args...)

	r.runCallbackForCmd(runCmd)

	fullCmd := strings.Join(runCmd, " ")

	results, found := r.processes[fullCmd]
	if !found {
		panic(fmt.Sprintf("Failed to find process for %s", fullCmd))
	}

	results[0].Stdout = cmd.Stdout
	results[0].Stderr = cmd.Stderr

	for _, proc := range results {
		if !proc.Waited {
			return proc, proc.StartErr
		}
	}

	panic(fmt.Sprintf("Failed to find available process for %s", fullCmd))
}

func (r *FakeCmdRunner) RunCommand(cmdName string, args ...string) (string, string, int, error) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	runCmd := append([]string{cmdName}, args...)
	r.RunCommands = append(r.RunCommands, runCmd)

	r.runCallbackForCmd(runCmd)

	return r.getOutputsForCmd(runCmd)
}

func (r *FakeCmdRunner) ClearCommandHistory() {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	r.RunCommands = [][]string{}
	r.RunCommandsQuietly = [][]string{}
	r.RunCommandsWithInput = [][]string{}
}

func (r *FakeCmdRunner) RunCommandQuietly(cmdName string, args ...string) (string, string, int, error) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	runCmd := append([]string{cmdName}, args...)
	r.RunCommandsQuietly = append(r.RunCommandsQuietly, runCmd)

	r.runCallbackForCmd(runCmd)

	return r.getOutputsForCmd(runCmd)
}

func (r *FakeCmdRunner) RunCommandWithInput(input, cmdName string, args ...string) (string, string, int, error) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	runCmd := append([]string{input, cmdName}, args...)
	r.RunCommandsWithInput = append(r.RunCommandsWithInput, runCmd)

	r.runCallbackForCmd(runCmd)

	return r.getOutputsForCmd(runCmd)
}

func (r *FakeCmdRunner) CommandExists(cmdName string) bool {
	return r.CommandExistsValue || r.AvailableCommands[cmdName]
}

func (r *FakeCmdRunner) AddCmdResult(fullCmd string, result FakeCmdResult) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	results := r.commandResults[fullCmd]
	r.commandResults[fullCmd] = append(results, result)
}

func (r *FakeCmdRunner) AddProcess(fullCmd string, process *FakeProcess) {
	r.processesLock.Lock()
	defer r.processesLock.Unlock()

	processes := r.processes[fullCmd]
	r.processes[fullCmd] = append(processes, process)
}

func (r *FakeCmdRunner) SetCmdCallback(fullCmd string, callback FakeCmdCallback) {
	r.runCommandCallbacks[fullCmd] = callback
}

func (r *FakeCmdRunner) getOutputsForCmd(runCmd []string) (string, string, int, error) {
	fullCmd := strings.Join(runCmd, " ")
	results, found := r.commandResults[fullCmd]
	if found {
		result := results[0]
		newResults := []FakeCmdResult{}
		if len(results) > 1 {
			newResults = results[1:]
		}

		if !result.Sticky {
			if len(newResults) == 0 {
				delete(r.commandResults, "fullCmd")
			} else {
				r.commandResults[fullCmd] = newResults
			}
		}

		return result.Stdout, result.Stderr, result.ExitStatus, result.Error
	}

	return "", "", -1, nil
}

func (r *FakeCmdRunner) runCallbackForCmd(runCmd []string) {
	fullCmd := strings.Join(runCmd, " ")
	runCmdCallback, found := r.runCommandCallbacks[fullCmd]
	if found {
		runCmdCallback()
	}
}
<commit_msg>Properly remove Cmd from commandResults<commit_after>package fakes

import (
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	boshsys "github.com/cloudfoundry/bosh-utils/system"
)

type FakeCmdRunner struct {
	commandResults     map[string][]FakeCmdResult
	commandResultsLock sync.Mutex

	processes     map[string][]*FakeProcess
	processesLock sync.Mutex

	RunComplexCommands   []boshsys.Command
	RunCommands          [][]string
	RunCommandsWithInput [][]string
	RunCommandsQuietly   [][]string
	runCommandCallbacks  map[string]FakeCmdCallback

	CommandExistsValue bool
	AvailableCommands  map[string]bool
}

type FakeCmdCallback func()

type FakeCmdResult struct {
	Stdout     string
	Stderr     string
	ExitStatus int
	Error      error

	Sticky bool // Set to true if this result should ALWAYS be returned for the given command
}

type FakeProcess struct {
	StartErr error

	WaitCh chan boshsys.Result

	Waited     bool
	WaitResult boshsys.Result

	TerminatedNicely               bool
	TerminatedNicelyCallBack       func(*FakeProcess)
	TerminateNicelyKillGracePeriod time.Duration
	TerminateNicelyErr             error

	Stdout io.Writer
	Stderr io.Writer
}

func (p *FakeProcess) Wait() <-chan boshsys.Result {
	if p.Waited {
		panic("Cannot Wait() on process multiple times")
	}

	p.Waited = true
	p.WaitCh = make(chan boshsys.Result, 1)

	if p.TerminatedNicelyCallBack == nil {
		p.WaitCh <- p.WaitResult
	}
	return p.WaitCh
}

func (p *FakeProcess) TerminateNicely(killGracePeriod time.Duration) error {
	p.TerminateNicelyKillGracePeriod = killGracePeriod
	p.TerminatedNicely = true
	if p.TerminatedNicelyCallBack != nil {
		p.TerminatedNicelyCallBack(p)
	}
	return p.TerminateNicelyErr
}

func NewFakeCmdRunner() *FakeCmdRunner {
	return &FakeCmdRunner{
		AvailableCommands:   map[string]bool{},
		commandResults:      map[string][]FakeCmdResult{},
		runCommandCallbacks: map[string]FakeCmdCallback{},
		processes:           map[string][]*FakeProcess{},
	}
}

func (r *FakeCmdRunner) RunComplexCommand(cmd boshsys.Command) (string, string, int, error) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	r.RunComplexCommands = append(r.RunComplexCommands, cmd)

	runCmd := append([]string{cmd.Name}, cmd.Args...)

	r.runCallbackForCmd(runCmd)

	stdout, stderr, exitstatus, err := r.getOutputsForCmd(runCmd)

	if cmd.Stdout != nil {
		cmd.Stdout.Write([]byte(stdout))
	}

	if cmd.Stderr != nil {
		cmd.Stderr.Write([]byte(stderr))
	}

	return stdout, stderr, exitstatus, err
}

func (r *FakeCmdRunner) RunComplexCommandAsync(cmd boshsys.Command) (boshsys.Process, error) {
	r.processesLock.Lock()
	defer r.processesLock.Unlock()

	r.RunComplexCommands = append(r.RunComplexCommands, cmd)

	runCmd := append([]string{cmd.Name}, cmd.Args...)

	r.runCallbackForCmd(runCmd)

	fullCmd := strings.Join(runCmd, " ")

	results, found := r.processes[fullCmd]
	if !found {
		panic(fmt.Sprintf("Failed to find process for %s", fullCmd))
	}

	results[0].Stdout = cmd.Stdout
	results[0].Stderr = cmd.Stderr

	for _, proc := range results {
		if !proc.Waited {
			return proc, proc.StartErr
		}
	}

	panic(fmt.Sprintf("Failed to find available process for %s", fullCmd))
}

func (r *FakeCmdRunner) RunCommand(cmdName string, args ...string) (string, string, int, error) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	runCmd := append([]string{cmdName}, args...)
	r.RunCommands = append(r.RunCommands, runCmd)

	r.runCallbackForCmd(runCmd)

	return r.getOutputsForCmd(runCmd)
}

func (r *FakeCmdRunner) ClearCommandHistory() {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	r.RunCommands = [][]string{}
	r.RunCommandsQuietly = [][]string{}
	r.RunCommandsWithInput = [][]string{}
}

func (r *FakeCmdRunner) RunCommandQuietly(cmdName string, args ...string) (string, string, int, error) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	runCmd := append([]string{cmdName}, args...)
	r.RunCommandsQuietly = append(r.RunCommandsQuietly, runCmd)

	r.runCallbackForCmd(runCmd)

	return r.getOutputsForCmd(runCmd)
}

func (r *FakeCmdRunner) RunCommandWithInput(input, cmdName string, args ...string) (string, string, int, error) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	runCmd := append([]string{input, cmdName}, args...)
	r.RunCommandsWithInput = append(r.RunCommandsWithInput, runCmd)

	r.runCallbackForCmd(runCmd)

	return r.getOutputsForCmd(runCmd)
}

func (r *FakeCmdRunner) CommandExists(cmdName string) bool {
	return r.CommandExistsValue || r.AvailableCommands[cmdName]
}

func (r *FakeCmdRunner) AddCmdResult(fullCmd string, result FakeCmdResult) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	results := r.commandResults[fullCmd]
	r.commandResults[fullCmd] = append(results, result)
}

func (r *FakeCmdRunner) AddProcess(fullCmd string, process *FakeProcess) {
	r.processesLock.Lock()
	defer r.processesLock.Unlock()

	processes := r.processes[fullCmd]
	r.processes[fullCmd] = append(processes, process)
}

func (r *FakeCmdRunner) SetCmdCallback(fullCmd string, callback FakeCmdCallback) {
	r.runCommandCallbacks[fullCmd] = callback
}

func (r *FakeCmdRunner) getOutputsForCmd(runCmd []string) (string, string, int, error) {
	fullCmd := strings.Join(runCmd, " ")
	results, found := r.commandResults[fullCmd]
	if found {
		result := results[0]
		newResults := []FakeCmdResult{}
		if len(results) > 1 {
			newResults = results[1:]
		}

		if !result.Sticky {
			if len(newResults) == 0 {
				delete(r.commandResults, fullCmd)
			} else {
				r.commandResults[fullCmd] = newResults
			}
		}

		return result.Stdout, result.Stderr, result.ExitStatus, result.Error
	}

	return "", "", -1, nil
}

func (r *FakeCmdRunner) runCallbackForCmd(runCmd []string) {
	fullCmd := strings.Join(runCmd, " ")
	runCmdCallback, found := r.runCommandCallbacks[fullCmd]
	if found {
		runCmdCallback()
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
	"text/template"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/rexray/gocsi/utils"
	"google.golang.org/grpc"
	"google.golang.org/grpc/status"
)

var debug, _ = strconv.ParseBool(os.Getenv("X_CSI_DEBUG"))

var root struct {
	ctx       context.Context
	client    *grpc.ClientConn
	tpl       *template.Template
	userCreds map[string]string

	genMarkdown bool
	logLevel    logLevelArg
	format      string
	endpoint    string
	insecure    bool
	timeout     time.Duration
	version     csiVersionArg
	metadata    mapOfStringArg

	withReqLogging bool
	withRepLogging bool

	withSpecValidator            bool
	withRequiresCreds            bool
	withRequiresNodeID           bool
	withRequiresPubVolInfo       bool
	withRequiresVolumeAttributes bool
}

var (
	activeArgs []string
	activeCmd  *cobra.Command
)

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "csc",
	Short: "a command line container storage interface (CSI) client",
	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {

		// Enable debug level logging and request and response logging
		// if the environment variable that controls deubg mode is set
		// to a truthy value.
		if debug {
			root.logLevel.Set(log.DebugLevel.String())
			root.withReqLogging = true
			root.withReqLogging = true
		}

		// Set the log level.
		lvl, _ := root.logLevel.Val()
		log.SetLevel(lvl)

		if debug {
			log.Warn("debug mode enabled")
		}

		root.ctx = context.Background()
		log.Debug("assigned the root context")

		// Initialize the template if necessary.
		if root.format == "" {
			switch cmd.Name() {
			case listVolumesCmd.Name():
				if listVolumes.paging {
					root.format = volumeInfoFormat
				} else {
					root.format = listVolumesFormat
				}
			case createVolumeCmd.Name():
				root.format = volumeInfoFormat
			case supportedVersCmd.Name():
				root.format = supportedVersionsFormat
			case pluginInfoCmd.Name():
				root.format = pluginInfoFormat
			case pluginCapsCmd.Name():
				root.format = pluginCapsFormat
			}
		}
		if root.format != "" {
			tpl, err := template.New("t").Funcs(template.FuncMap{
				"isa": func(o interface{}, t string) bool {
					return fmt.Sprintf("%T", o) == t
				},
			}).Parse(root.format)
			if err != nil {
				return err
			}
			root.tpl = tpl
		}

		// Parse the credentials if they exist.
		root.userCreds = utils.ParseMap(os.Getenv("X_CSI_USER_CREDENTIALS"))

		// Create the gRPC client connection.
		opts := []grpc.DialOption{
			grpc.WithDialer(
				func(target string, timeout time.Duration) (net.Conn, error) {
					proto, addr, err := utils.ParseProtoAddr(target)
					if err != nil {
						return nil, err
					}
					return net.DialTimeout(proto, addr, timeout)
				}),
		}

		// Disable TLS if specified.
		if root.insecure {
			opts = append(opts, grpc.WithInsecure())
		}

		// Add interceptors to the client if any are configured.
		if o := getClientInterceptorsDialOpt(); o != nil {
			opts = append(opts, o)
		}

		ctx, cancel := context.WithTimeout(root.ctx, root.timeout)
		defer cancel()
		client, err := grpc.DialContext(ctx, root.endpoint, opts...)
		if err != nil {
			return err
		}
		root.client = client

		return nil
	},
}

// Execute adds all child commands to the root command and sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	if err := RootCmd.Execute(); err != nil {
		exitCode := 1
		if stat, ok := status.FromError(err); ok {
			exitCode = int(stat.Code())
			fmt.Fprintln(os.Stderr, stat.Message())
		} else {
			fmt.Fprintf(os.Stderr, "%v\n", err)
		}
		fmt.Fprintf(os.Stderr, "\nPlease use -h,--help for more information\n")
		os.Exit(exitCode)
	}
}

func init() {
	setHelpAndUsage(RootCmd)

	flagLogLevel(
		RootCmd.PersistentFlags(),
		&root.logLevel,
		"warn")

	flagEndpoint(
		RootCmd.PersistentFlags(),
		&root.endpoint,
		os.Getenv("CSI_ENDPOINT"))

	flagTimeout(
		RootCmd.PersistentFlags(),
		&root.timeout,
		"1m")

	flagWithRequestLogging(
		RootCmd.PersistentFlags(),
		&root.withReqLogging,
		"false")

	flagWithResponseLogging(
		RootCmd.PersistentFlags(),
		&root.withRepLogging,
		"false")

	flagWithSpecValidation(
		RootCmd.PersistentFlags(),
		&root.withSpecValidator,
		"false")

	RootCmd.PersistentFlags().BoolVarP(
		&root.insecure,
		"insecure",
		"i",
		true,
		`Disables transport security for the client via the gRPC dial option
        WithInsecure (https://goo.gl/Y95SfW)`)

	RootCmd.PersistentFlags().VarP(
		&root.metadata,
		"metadata",
		"m",
		`Sets one or more key/value pairs to use as gRPC metadata sent with all
        RPCs. gRPC metadata is similar to HTTP headers. For example:

            --metadata key1=val1 --m key2=val2,key3=val3

            -m key1=val1,key2=val2 --metadata key3=val3

        Read more on gRPC metadata at https://goo.gl/iTci67`)

	RootCmd.PersistentFlags().VarP(
		&root.version,
		"version",
		"v",
		`The version sent with an RPC may be specified as MAJOR.MINOR.PATCH`)

}

type logger struct {
	f func(msg string, args ...interface{})
	w io.Writer
}

func newLogger(f func(msg string, args ...interface{})) *logger {
	l := &logger{f: f}
	r, w := io.Pipe()
	l.w = w
	go func() {
		scan := bufio.NewScanner(r)
		for scan.Scan() {
			f(scan.Text())
		}
	}()
	return l
}

func (l *logger) Write(data []byte) (int, error) {
	return l.w.Write(data)
}
<commit_msg>Fix TCP endpoint support for csc<commit_after>package cmd

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
	"text/template"
	"time"

	"github.com/rexray/gocsi/utils"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"google.golang.org/grpc"
	"google.golang.org/grpc/status"
)

var debug, _ = strconv.ParseBool(os.Getenv("X_CSI_DEBUG"))

var root struct {
	ctx       context.Context
	client    *grpc.ClientConn
	tpl       *template.Template
	userCreds map[string]string

	genMarkdown bool
	logLevel    logLevelArg
	format      string
	endpoint    string
	insecure    bool
	timeout     time.Duration
	version     csiVersionArg
	metadata    mapOfStringArg

	withReqLogging bool
	withRepLogging bool

	withSpecValidator            bool
	withRequiresCreds            bool
	withRequiresNodeID           bool
	withRequiresPubVolInfo       bool
	withRequiresVolumeAttributes bool
}

var (
	activeArgs []string
	activeCmd  *cobra.Command
)

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "csc",
	Short: "a command line container storage interface (CSI) client",
	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {

		// Enable debug level logging and request and response logging
		// if the environment variable that controls deubg mode is set
		// to a truthy value.
		if debug {
			root.logLevel.Set(log.DebugLevel.String())
			root.withReqLogging = true
			root.withReqLogging = true
		}

		// Set the log level.
		lvl, _ := root.logLevel.Val()
		log.SetLevel(lvl)

		if debug {
			log.Warn("debug mode enabled")
		}

		root.ctx = context.Background()
		log.Debug("assigned the root context")

		// Initialize the template if necessary.
		if root.format == "" {
			switch cmd.Name() {
			case listVolumesCmd.Name():
				if listVolumes.paging {
					root.format = volumeInfoFormat
				} else {
					root.format = listVolumesFormat
				}
			case createVolumeCmd.Name():
				root.format = volumeInfoFormat
			case supportedVersCmd.Name():
				root.format = supportedVersionsFormat
			case pluginInfoCmd.Name():
				root.format = pluginInfoFormat
			case pluginCapsCmd.Name():
				root.format = pluginCapsFormat
			}
		}
		if root.format != "" {
			tpl, err := template.New("t").Funcs(template.FuncMap{
				"isa": func(o interface{}, t string) bool {
					return fmt.Sprintf("%T", o) == t
				},
			}).Parse(root.format)
			if err != nil {
				return err
			}
			root.tpl = tpl
		}

		// Parse the credentials if they exist.
		root.userCreds = utils.ParseMap(os.Getenv("X_CSI_USER_CREDENTIALS"))

		// Create the gRPC client connection.
		opts := []grpc.DialOption{
			grpc.WithDialer(
				func(string, time.Duration) (net.Conn, error) {
					proto, addr, err := utils.ParseProtoAddr(root.endpoint)
					log.WithFields(map[string]interface{}{
						"proto":   proto,
						"addr":    addr,
						"timeout": root.timeout,
					}).Debug("parsed endpoint info")
					if err != nil {
						return nil, err
					}
					return net.DialTimeout(proto, addr, root.timeout)
				}),
		}

		// Disable TLS if specified.
		if root.insecure {
			opts = append(opts, grpc.WithInsecure())
		}

		// Add interceptors to the client if any are configured.
		if o := getClientInterceptorsDialOpt(); o != nil {
			opts = append(opts, o)
		}

		ctx, cancel := context.WithTimeout(root.ctx, root.timeout)
		defer cancel()
		client, err := grpc.DialContext(ctx, root.endpoint, opts...)
		if err != nil {
			return err
		}
		root.client = client

		return nil
	},
}

// Execute adds all child commands to the root command and sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	if err := RootCmd.Execute(); err != nil {
		exitCode := 1
		if stat, ok := status.FromError(err); ok {
			exitCode = int(stat.Code())
			fmt.Fprintln(os.Stderr, stat.Message())
		} else {
			fmt.Fprintf(os.Stderr, "%v\n", err)
		}
		fmt.Fprintf(os.Stderr, "\nPlease use -h,--help for more information\n")
		os.Exit(exitCode)
	}
}

func init() {
	setHelpAndUsage(RootCmd)

	flagLogLevel(
		RootCmd.PersistentFlags(),
		&root.logLevel,
		"warn")

	flagEndpoint(
		RootCmd.PersistentFlags(),
		&root.endpoint,
		os.Getenv("CSI_ENDPOINT"))

	flagTimeout(
		RootCmd.PersistentFlags(),
		&root.timeout,
		"1m")

	flagWithRequestLogging(
		RootCmd.PersistentFlags(),
		&root.withReqLogging,
		"false")

	flagWithResponseLogging(
		RootCmd.PersistentFlags(),
		&root.withRepLogging,
		"false")

	flagWithSpecValidation(
		RootCmd.PersistentFlags(),
		&root.withSpecValidator,
		"false")

	RootCmd.PersistentFlags().BoolVarP(
		&root.insecure,
		"insecure",
		"i",
		true,
		`Disables transport security for the client via the gRPC dial option
        WithInsecure (https://goo.gl/Y95SfW)`)

	RootCmd.PersistentFlags().VarP(
		&root.metadata,
		"metadata",
		"m",
		`Sets one or more key/value pairs to use as gRPC metadata sent with all
        RPCs. gRPC metadata is similar to HTTP headers. For example:

            --metadata key1=val1 --m key2=val2,key3=val3

            -m key1=val1,key2=val2 --metadata key3=val3

        Read more on gRPC metadata at https://goo.gl/iTci67`)

	RootCmd.PersistentFlags().VarP(
		&root.version,
		"version",
		"v",
		`The version sent with an RPC may be specified as MAJOR.MINOR.PATCH`)

}

type logger struct {
	f func(msg string, args ...interface{})
	w io.Writer
}

func newLogger(f func(msg string, args ...interface{})) *logger {
	l := &logger{f: f}
	r, w := io.Pipe()
	l.w = w
	go func() {
		scan := bufio.NewScanner(r)
		for scan.Scan() {
			f(scan.Text())
		}
	}()
	return l
}

func (l *logger) Write(data []byte) (int, error) {
	return l.w.Write(data)
}
<|endoftext|>


<commit_before>package assets

import (
	"fmt"
	"io"
	"strconv"

	"github.com/ugorji/go/codec"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/unversioned"
)

var (
	suite              = "pachyderm"
	pachdImage         = "pachyderm/pachd"
	etcdImage          = "gcr.io/google_containers/etcd:2.0.12"
	rethinkImage       = "rethinkdb:2.1.5"
	serviceAccountName = "pachyderm"
	etcdName           = "etcd"
	pachdName          = "pachd"
	rethinkName        = "rethink"
	amazonSecretName   = "amazon-secret"
	googleSecretName   = "google-secret"
	trueVal            = true
)

type backend int

const (
	localBackend backend = iota
	amazonBackend
	googleBackend
)

func ServiceAccount() *api.ServiceAccount {
	return &api.ServiceAccount{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "ServiceAccount",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   serviceAccountName,
			Labels: labels(""),
		},
	}
}

//PachdRc TODO secrets is only necessary because dockerized kube chokes on them
func PachdRc(shards uint64, backend backend) *api.ReplicationController {
	volumes := []api.Volume{
		{
			Name: "pach-disk",
		},
	}
	volumeMounts := []api.VolumeMount{
		{
			Name:      "pach-disk",
			MountPath: "/pach",
		},
	}
	switch backend {
	case localBackend:
	case amazonBackend:
		volumes = append(volumes, api.Volume{
			Name: amazonSecretName,
			VolumeSource: api.VolumeSource{
				Secret: &api.SecretVolumeSource{
					SecretName: amazonSecretName,
				},
			},
		})
		volumeMounts = append(volumeMounts, api.VolumeMount{
			Name:      amazonSecretName,
			MountPath: "/" + amazonSecretName,
		})
	case googleBackend:
		volumes = append(volumes, api.Volume{
			Name: googleSecretName,
			VolumeSource: api.VolumeSource{
				Secret: &api.SecretVolumeSource{
					SecretName: googleSecretName,
				},
			},
		})
		volumeMounts = append(volumeMounts, api.VolumeMount{
			Name:      googleSecretName,
			MountPath: "/" + googleSecretName,
		})
	}
	return &api.ReplicationController{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "ReplicationController",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   pachdName,
			Labels: labels(pachdName),
		},
		Spec: api.ReplicationControllerSpec{
			Replicas: 1,
			Selector: map[string]string{
				"app": pachdName,
			},
			Template: &api.PodTemplateSpec{
				ObjectMeta: api.ObjectMeta{
					Name:   pachdName,
					Labels: labels(pachdName),
				},
				Spec: api.PodSpec{
					Containers: []api.Container{
						{
							Name:  pachdName,
							Image: pachdImage,
							Env: []api.EnvVar{
								{
									Name:  "PACH_ROOT",
									Value: "/pach",
								},
								{
									Name:  "NUM_SHARDS",
									Value: strconv.FormatUint(shards, 10),
								},
							},
							Ports: []api.ContainerPort{
								{
									ContainerPort: 650,
									Protocol:      "TCP",
									Name:          "api-grpc-port",
								},
								{
									ContainerPort: 750,
									Name:          "api-http-port",
								},
								{
									ContainerPort: 1050,
									Name:          "trace-port",
								},
							},
							VolumeMounts: volumeMounts,
							SecurityContext: &api.SecurityContext{
								Privileged: &trueVal, // god is this dumb
							},
							ImagePullPolicy: "IfNotPresent",
						},
					},
					ServiceAccountName: serviceAccountName,
					Volumes:            volumes,
				},
			},
		},
	}
}

func PachdService() *api.Service {
	return &api.Service{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "Service",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   pachdName,
			Labels: labels(pachdName),
		},
		Spec: api.ServiceSpec{
			Type: api.ServiceTypeNodePort,
			Selector: map[string]string{
				"app": pachdName,
			},
			Ports: []api.ServicePort{
				{
					Port:     650,
					Name:     "api-grpc-port",
					NodePort: 30650,
				},
				{
					Port:     750,
					Name:     "api-http-port",
					NodePort: 30750,
				},
			},
		},
	}
}

func EtcdRc() *api.ReplicationController {
	return &api.ReplicationController{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "ReplicationController",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   etcdName,
			Labels: labels(etcdName),
		},
		Spec: api.ReplicationControllerSpec{
			Replicas: 1,
			Selector: map[string]string{
				"app": etcdName,
			},
			Template: &api.PodTemplateSpec{
				ObjectMeta: api.ObjectMeta{
					Name:   etcdName,
					Labels: labels(etcdName),
				},
				Spec: api.PodSpec{
					Containers: []api.Container{
						{
							Name:  etcdName,
							Image: etcdImage,
							//TODO figure out how to get a cluster of these to talk to each other
							Command: []string{"/usr/local/bin/etcd", "--bind-addr=0.0.0.0:2379", "--data-dir=/var/etcd/data"},
							Ports: []api.ContainerPort{
								{
									ContainerPort: 2379,
									Name:          "client-port",
								},
								{
									ContainerPort: 2380,
									Name:          "peer-port",
								},
							},
							VolumeMounts: []api.VolumeMount{
								{
									Name:      "etcd-storage",
									MountPath: "/var/data/etcd",
								},
							},
							ImagePullPolicy: "IfNotPresent",
						},
					},
					Volumes: []api.Volume{
						{
							Name: "etcd-storage",
						},
					},
				},
			},
		},
	}
}

func EtcdService() *api.Service {
	return &api.Service{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "Service",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   etcdName,
			Labels: labels(etcdName),
		},
		Spec: api.ServiceSpec{
			Selector: map[string]string{
				"app": etcdName,
			},
			Ports: []api.ServicePort{
				{
					Port: 2379,
					Name: "client-port",
				},
				{
					Port: 2380,
					Name: "peer-port",
				},
			},
		},
	}
}

func RethinkRc() *api.ReplicationController {
	return &api.ReplicationController{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "ReplicationController",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   rethinkName,
			Labels: labels(rethinkName),
		},
		Spec: api.ReplicationControllerSpec{
			Replicas: 1,
			Selector: map[string]string{
				"app": rethinkName,
			},
			Template: &api.PodTemplateSpec{
				ObjectMeta: api.ObjectMeta{
					Name:   rethinkName,
					Labels: labels(rethinkName),
				},
				Spec: api.PodSpec{
					Containers: []api.Container{
						{
							Name:  rethinkName,
							Image: rethinkImage,
							//TODO figure out how to get a cluster of these to talk to each other
							Command: []string{"rethinkdb", "-d", "/var/rethinkdb/data", "--bind", "all"},
							Ports: []api.ContainerPort{
								{
									ContainerPort: 8080,
									Name:          "admin-port",
								},
								{
									ContainerPort: 28015,
									Name:          "driver-port",
								},
								{
									ContainerPort: 29015,
									Name:          "cluster-port",
								},
							},
							VolumeMounts: []api.VolumeMount{
								{
									Name:      "rethink-storage",
									MountPath: "/var/rethinkdb/data",
								},
							},
							ImagePullPolicy: "IfNotPresent",
						},
					},
					Volumes: []api.Volume{
						{
							Name: "rethink-storage",
						},
						//TODO this needs to be real storage
					},
				},
			},
		},
	}
}

func RethinkService() *api.Service {
	return &api.Service{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "Service",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   rethinkName,
			Labels: labels(rethinkName),
		},
		Spec: api.ServiceSpec{
			Selector: map[string]string{
				"app": rethinkName,
			},
			Ports: []api.ServicePort{
				{
					Port: 8080,
					Name: "admin-port",
				},
				{
					Port: 28015,
					Name: "driver-port",
				},
				{
					Port: 29015,
					Name: "cluster-port",
				},
			},
		},
	}
}

func AmazonSecret(bucket string, id string, secret string, token string, region string) *api.Secret {
	return &api.Secret{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "Secret",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   amazonSecretName,
			Labels: labels(amazonSecretName),
		},
		Data: map[string][]byte{
			"bucket": []byte(bucket),
			"id":     []byte(id),
			"secret": []byte(secret),
			"token":  []byte(token),
			"region": []byte(region),
		},
	}
}

func GoogleSecret(bucket string) *api.Secret {
	return &api.Secret{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "Secret",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   googleSecretName,
			Labels: labels(googleSecretName),
		},
		Data: map[string][]byte{
			"bucket": []byte(bucket),
		},
	}
}

// WriteAssets creates the assets in a dir. It expects dir to already exist.
func WriteAssets(w io.Writer, shards uint64, backend backend) {
	encoder := codec.NewEncoder(w, &codec.JsonHandle{Indent: 2})

	ServiceAccount().CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")

	EtcdRc().CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")
	EtcdService().CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")

	RethinkService().CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")
	RethinkRc().CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")

	PachdService().CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")
	PachdRc(shards, backend).CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")

}

func WriteLocalAssets(w io.Writer, shards uint64) {
	WriteAssets(w, shards, localBackend)
}

func WriteAmazonAssets(w io.Writer, shards uint64, bucket string, id string, secret string, token string, region string) {
	WriteAssets(w, shards, amazonBackend)
	encoder := codec.NewEncoder(w, &codec.JsonHandle{Indent: 2})
	AmazonSecret(bucket, id, secret, token, region).CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")
}

func WriteGoogleAssets(w io.Writer, shards uint64, bucket string) {
	WriteAssets(w, shards, googleBackend)
	encoder := codec.NewEncoder(w, &codec.JsonHandle{Indent: 2})
	GoogleSecret(bucket).CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")
}

func labels(name string) map[string]string {
	return map[string]string{
		"app":   name,
		"suite": suite,
	}
}
<commit_msg>Make default replicas 8.<commit_after>package assets

import (
	"fmt"
	"io"
	"strconv"

	"github.com/ugorji/go/codec"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/unversioned"
)

var (
	suite              = "pachyderm"
	pachdImage         = "pachyderm/pachd"
	etcdImage          = "gcr.io/google_containers/etcd:2.0.12"
	rethinkImage       = "rethinkdb:2.1.5"
	serviceAccountName = "pachyderm"
	etcdName           = "etcd"
	pachdName          = "pachd"
	rethinkName        = "rethink"
	amazonSecretName   = "amazon-secret"
	googleSecretName   = "google-secret"
	trueVal            = true
)

type backend int

const (
	localBackend backend = iota
	amazonBackend
	googleBackend
)

func ServiceAccount() *api.ServiceAccount {
	return &api.ServiceAccount{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "ServiceAccount",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   serviceAccountName,
			Labels: labels(""),
		},
	}
}

//PachdRc TODO secrets is only necessary because dockerized kube chokes on them
func PachdRc(shards uint64, backend backend) *api.ReplicationController {
	volumes := []api.Volume{
		{
			Name: "pach-disk",
		},
	}
	volumeMounts := []api.VolumeMount{
		{
			Name:      "pach-disk",
			MountPath: "/pach",
		},
	}
	switch backend {
	case localBackend:
	case amazonBackend:
		volumes = append(volumes, api.Volume{
			Name: amazonSecretName,
			VolumeSource: api.VolumeSource{
				Secret: &api.SecretVolumeSource{
					SecretName: amazonSecretName,
				},
			},
		})
		volumeMounts = append(volumeMounts, api.VolumeMount{
			Name:      amazonSecretName,
			MountPath: "/" + amazonSecretName,
		})
	case googleBackend:
		volumes = append(volumes, api.Volume{
			Name: googleSecretName,
			VolumeSource: api.VolumeSource{
				Secret: &api.SecretVolumeSource{
					SecretName: googleSecretName,
				},
			},
		})
		volumeMounts = append(volumeMounts, api.VolumeMount{
			Name:      googleSecretName,
			MountPath: "/" + googleSecretName,
		})
	}
	return &api.ReplicationController{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "ReplicationController",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   pachdName,
			Labels: labels(pachdName),
		},
		Spec: api.ReplicationControllerSpec{
			Replicas: 8,
			Selector: map[string]string{
				"app": pachdName,
			},
			Template: &api.PodTemplateSpec{
				ObjectMeta: api.ObjectMeta{
					Name:   pachdName,
					Labels: labels(pachdName),
				},
				Spec: api.PodSpec{
					Containers: []api.Container{
						{
							Name:  pachdName,
							Image: pachdImage,
							Env: []api.EnvVar{
								{
									Name:  "PACH_ROOT",
									Value: "/pach",
								},
								{
									Name:  "NUM_SHARDS",
									Value: strconv.FormatUint(shards, 10),
								},
							},
							Ports: []api.ContainerPort{
								{
									ContainerPort: 650,
									Protocol:      "TCP",
									Name:          "api-grpc-port",
								},
								{
									ContainerPort: 750,
									Name:          "api-http-port",
								},
								{
									ContainerPort: 1050,
									Name:          "trace-port",
								},
							},
							VolumeMounts: volumeMounts,
							SecurityContext: &api.SecurityContext{
								Privileged: &trueVal, // god is this dumb
							},
							ImagePullPolicy: "IfNotPresent",
						},
					},
					ServiceAccountName: serviceAccountName,
					Volumes:            volumes,
				},
			},
		},
	}
}

func PachdService() *api.Service {
	return &api.Service{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "Service",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   pachdName,
			Labels: labels(pachdName),
		},
		Spec: api.ServiceSpec{
			Type: api.ServiceTypeNodePort,
			Selector: map[string]string{
				"app": pachdName,
			},
			Ports: []api.ServicePort{
				{
					Port:     650,
					Name:     "api-grpc-port",
					NodePort: 30650,
				},
				{
					Port:     750,
					Name:     "api-http-port",
					NodePort: 30750,
				},
			},
		},
	}
}

func EtcdRc() *api.ReplicationController {
	return &api.ReplicationController{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "ReplicationController",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   etcdName,
			Labels: labels(etcdName),
		},
		Spec: api.ReplicationControllerSpec{
			Replicas: 1,
			Selector: map[string]string{
				"app": etcdName,
			},
			Template: &api.PodTemplateSpec{
				ObjectMeta: api.ObjectMeta{
					Name:   etcdName,
					Labels: labels(etcdName),
				},
				Spec: api.PodSpec{
					Containers: []api.Container{
						{
							Name:  etcdName,
							Image: etcdImage,
							//TODO figure out how to get a cluster of these to talk to each other
							Command: []string{"/usr/local/bin/etcd", "--bind-addr=0.0.0.0:2379", "--data-dir=/var/etcd/data"},
							Ports: []api.ContainerPort{
								{
									ContainerPort: 2379,
									Name:          "client-port",
								},
								{
									ContainerPort: 2380,
									Name:          "peer-port",
								},
							},
							VolumeMounts: []api.VolumeMount{
								{
									Name:      "etcd-storage",
									MountPath: "/var/data/etcd",
								},
							},
							ImagePullPolicy: "IfNotPresent",
						},
					},
					Volumes: []api.Volume{
						{
							Name: "etcd-storage",
						},
					},
				},
			},
		},
	}
}

func EtcdService() *api.Service {
	return &api.Service{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "Service",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   etcdName,
			Labels: labels(etcdName),
		},
		Spec: api.ServiceSpec{
			Selector: map[string]string{
				"app": etcdName,
			},
			Ports: []api.ServicePort{
				{
					Port: 2379,
					Name: "client-port",
				},
				{
					Port: 2380,
					Name: "peer-port",
				},
			},
		},
	}
}

func RethinkRc() *api.ReplicationController {
	return &api.ReplicationController{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "ReplicationController",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   rethinkName,
			Labels: labels(rethinkName),
		},
		Spec: api.ReplicationControllerSpec{
			Replicas: 1,
			Selector: map[string]string{
				"app": rethinkName,
			},
			Template: &api.PodTemplateSpec{
				ObjectMeta: api.ObjectMeta{
					Name:   rethinkName,
					Labels: labels(rethinkName),
				},
				Spec: api.PodSpec{
					Containers: []api.Container{
						{
							Name:  rethinkName,
							Image: rethinkImage,
							//TODO figure out how to get a cluster of these to talk to each other
							Command: []string{"rethinkdb", "-d", "/var/rethinkdb/data", "--bind", "all"},
							Ports: []api.ContainerPort{
								{
									ContainerPort: 8080,
									Name:          "admin-port",
								},
								{
									ContainerPort: 28015,
									Name:          "driver-port",
								},
								{
									ContainerPort: 29015,
									Name:          "cluster-port",
								},
							},
							VolumeMounts: []api.VolumeMount{
								{
									Name:      "rethink-storage",
									MountPath: "/var/rethinkdb/data",
								},
							},
							ImagePullPolicy: "IfNotPresent",
						},
					},
					Volumes: []api.Volume{
						{
							Name: "rethink-storage",
						},
						//TODO this needs to be real storage
					},
				},
			},
		},
	}
}

func RethinkService() *api.Service {
	return &api.Service{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "Service",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   rethinkName,
			Labels: labels(rethinkName),
		},
		Spec: api.ServiceSpec{
			Selector: map[string]string{
				"app": rethinkName,
			},
			Ports: []api.ServicePort{
				{
					Port: 8080,
					Name: "admin-port",
				},
				{
					Port: 28015,
					Name: "driver-port",
				},
				{
					Port: 29015,
					Name: "cluster-port",
				},
			},
		},
	}
}

func AmazonSecret(bucket string, id string, secret string, token string, region string) *api.Secret {
	return &api.Secret{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "Secret",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   amazonSecretName,
			Labels: labels(amazonSecretName),
		},
		Data: map[string][]byte{
			"bucket": []byte(bucket),
			"id":     []byte(id),
			"secret": []byte(secret),
			"token":  []byte(token),
			"region": []byte(region),
		},
	}
}

func GoogleSecret(bucket string) *api.Secret {
	return &api.Secret{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "Secret",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   googleSecretName,
			Labels: labels(googleSecretName),
		},
		Data: map[string][]byte{
			"bucket": []byte(bucket),
		},
	}
}

// WriteAssets creates the assets in a dir. It expects dir to already exist.
func WriteAssets(w io.Writer, shards uint64, backend backend) {
	encoder := codec.NewEncoder(w, &codec.JsonHandle{Indent: 2})

	ServiceAccount().CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")

	EtcdRc().CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")
	EtcdService().CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")

	RethinkService().CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")
	RethinkRc().CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")

	PachdService().CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")
	PachdRc(shards, backend).CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")

}

func WriteLocalAssets(w io.Writer, shards uint64) {
	WriteAssets(w, shards, localBackend)
}

func WriteAmazonAssets(w io.Writer, shards uint64, bucket string, id string, secret string, token string, region string) {
	WriteAssets(w, shards, amazonBackend)
	encoder := codec.NewEncoder(w, &codec.JsonHandle{Indent: 2})
	AmazonSecret(bucket, id, secret, token, region).CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")
}

func WriteGoogleAssets(w io.Writer, shards uint64, bucket string) {
	WriteAssets(w, shards, googleBackend)
	encoder := codec.NewEncoder(w, &codec.JsonHandle{Indent: 2})
	GoogleSecret(bucket).CodecEncodeSelf(encoder)
	fmt.Fprintf(w, "\n")
}

func labels(name string) map[string]string {
	return map[string]string{
		"app":   name,
		"suite": suite,
	}
}
<|endoftext|>


<commit_before>package python

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"

	"strings"

	"sourcegraph.com/sourcegraph/srclib/graph"
	"sourcegraph.com/sourcegraph/srclib/toolchain"
	"sourcegraph.com/sourcegraph/srclib/unit"
)

type GraphContext struct {
	Unit *unit.SourceUnit
	Reqs []*requirement
}

func NewGraphContext(unit *unit.SourceUnit) *GraphContext {
	var g GraphContext
	g.Unit = unit
	for _, dep := range unit.Dependencies {
		if req, err := asRequirement(dep); err == nil {
			g.Reqs = append(g.Reqs, req)
		}
	}
	return &g
}

// Graphs the Python source unit. If run outside of a Docker container, this assumes that the source unit has already
// been installed (via pip or `python setup.py install`).
func (c *GraphContext) Graph() (*graph.Output, error) {
	programMode := os.Getenv("IN_DOCKER_CONTAINER") == ""
	tc, err := toolchain.Lookup("sourcegraph.com/sourcegraph/srclib-python")
	if err != nil {
		return nil, err
	}
	pipBin := "pip"
	pythonBin := "python"

	if programMode {
		tempPath, err := getTempPath()
		if err != nil {
			return nil, err
		}
		envName := fmt.Sprintf("%s-%s-env", getHash(c.Unit.Dir), url.QueryEscape(c.Unit.Name))
		envDir := filepath.Join(tempPath, envName)

		if _, err := os.Stat(filepath.Join(envDir)); os.IsNotExist(err) {
			// We don't have virtual env for this SourceUnit, create one.
			tcVENVBinPath := filepath.Join(tc.Dir, ".env", "bin")
			cmd := exec.Command(filepath.Join(tcVENVBinPath, "virtualenv"), envDir)
			if err := runCmdStderr(cmd); err != nil {
				return nil, err
			}
		}
		// Use binaries from our virutal env.
		pipBin = filepath.Join(envDir, "bin", "pip")
		pythonBin = filepath.Join(envDir, "bin", "python")
	}

	// NOTE: this may cause an error when graphing any source unit that depends
	// on jedi (or any other dependency of the graph code)
	requirementFiles, err := filepath.Glob(filepath.Join(c.Unit.Dir, "*requirements*.txt"))
	if err != nil {
		return nil, err
	}
	if _, err := os.Stat(filepath.Join(c.Unit.Dir, "setup.py")); !os.IsNotExist(err) {
		runCmdLogError(exec.Command(pipBin, "install", "-I", c.Unit.Dir))
	}
	installPipRequirements(pipBin, requirementFiles)

	if programMode {
		// Unlike in docker mode, this environment doesn't have toolchain requirements installed,
		// so install them.
		// NOTE: Doing this last to ensure toolchain has priority when it comes to seting dependency
		//       versions.
		// Todo(MaikuMori): Use symlinks from toolchains virtualenv to project virtual env.
		requirementFile := filepath.Join(tc.Dir, "requirements.txt")
		if err := runCmdStderr(exec.Command(pipBin, "install", "-r", requirementFile)); err != nil {
			return nil, err
		}
		if err := runCmdStderr(exec.Command(pipBin, "install", "-I", tc.Dir)); err != nil {
			return nil, err
		}
	}

	cmd := exec.Command(pythonBin, "-m", "grapher.graph", "--verbose", "--dir", c.Unit.Dir, "--files")
	cmd.Args = append(cmd.Args, c.Unit.Files...)
	cmd.Stderr = os.Stderr
	log.Printf("Running %v", cmd.Args)
	b, err := cmd.Output()
	if err != nil {
		return nil, err
	}

	var raw RawOutput
	if err := json.Unmarshal(b, &raw); err != nil {
		return nil, err
	}

	out := c.transform(&raw, c.Unit)
	return out, nil
}

func (c *GraphContext) transform(raw *RawOutput, unit *unit.SourceUnit) *graph.Output {
	var out graph.Output

	for _, def := range raw.Defs {
		out.Defs = append(out.Defs, c.transformDef(def))
		if doc := c.transformDefDoc(def); doc != nil {
			out.Docs = append(out.Docs, doc)
		}
	}
	for _, ref := range raw.Refs {
		if outRef, err := c.transformRef(ref); err == nil {
			out.Refs = append(out.Refs, outRef)
		} else {
			log.Printf("Could not transform ref %v: %s", ref, err)
		}
	}

	return &out
}

var jediKindToDefKind = map[string]string{
	"statement":        "var",
	"statementelement": "var",
	"param":            "var",
	"module":           "module",
	"submodule":        "module",
	"class":            "type",
	"function":         "func",
	"lambda":           "func",
	"import":           "var",
}

func (c *GraphContext) transformDef(rawDef *RawDef) *graph.Def {
	return &graph.Def{
		DefKey: graph.DefKey{
			Repo:     c.Unit.Repo,
			Unit:     c.Unit.Name,
			UnitType: c.Unit.Type,
			Path:     string(rawDef.Path),
		},
		TreePath: string(rawDef.Path), // TODO: make this consistent w/ old way
		Kind:     jediKindToDefKind[rawDef.Kind],
		Name:     rawDef.Name,
		File:     rawDef.File,
		DefStart: rawDef.DefStart,
		DefEnd:   rawDef.DefEnd,
		Exported: rawDef.Exported,
		Data:     nil, // TODO
	}
}

func (c *GraphContext) transformRef(rawRef *RawRef) (*graph.Ref, error) {
	defUnit, err := c.inferSourceUnit(rawRef, c.Reqs)
	if err != nil {
		return nil, err
	}

	defPath := string(rawRef.DefPath)
	if defPath == "" {
		defPath = "."
	}

	return &graph.Ref{
		DefRepo:     defUnit.Repo,
		DefUnitType: defUnit.Type,
		DefUnit:     defUnit.Name,
		DefPath:     defPath,

		Repo:     c.Unit.Repo,
		Unit:     c.Unit.Name,
		UnitType: c.Unit.Type,

		File:  rawRef.File,
		Start: rawRef.Start,
		End:   rawRef.End,
		Def:   rawRef.Def,
	}, nil
}

func (c *GraphContext) transformDefDoc(rawDef *RawDef) *graph.Doc {
	return nil
}

func (c *GraphContext) inferSourceUnit(rawRef *RawRef, reqs []*requirement) (*unit.SourceUnit, error) {
	if rawRef.ToBuiltin {
		return stdLibPkg.SourceUnit(), nil
	}
	return c.inferSourceUnitFromFile(rawRef.DefFile, reqs)
}

// Note: file is expected to be an absolute path
func (c *GraphContext) inferSourceUnitFromFile(file string, reqs []*requirement) (*unit.SourceUnit, error) {
	// Case: in current source unit (u)
	pwd, _ := os.Getwd()
	if isSubPath(pwd, file) {
		return c.Unit, nil
	}

	// Case: in dependent source unit(depUnits)
	fileCmps := strings.Split(file, string(filepath.Separator))
	pkgsDirIdx := -1
	for i, cmp := range fileCmps {
		if cmp == "site-packages" || cmp == "dist-packages" {
			pkgsDirIdx = i
			break
		}
	}
	if pkgsDirIdx != -1 {
		fileSubCmps := fileCmps[pkgsDirIdx+1:]
		fileSubPath := filepath.Join(fileSubCmps...)

		var foundReq *requirement
	FindReq:
		for _, req := range reqs {
			for _, pkg := range req.Packages {
				if isSubPath(moduleToFilepath(pkg, true), fileSubPath) {
					foundReq = req
					break FindReq
				}
			}
			for _, mod := range req.Modules {
				if moduleToFilepath(mod, false) == fileSubPath {
					foundReq = req
					break FindReq
				}
			}
		}

		if foundReq == nil {
			var candidatesStr string
			if len(reqs) <= 7 {
				candidatesStr = fmt.Sprintf("%v", reqs)
			} else {
				candidatesStr = fmt.Sprintf("%v...", reqs[:7])
			}
			// FIXME: This doesn't work, note the pointer in `[]*requirement`. As error you get
			//        string representation of array of pointers.
			return nil, fmt.Errorf("Could not find requirement that contains file %s. Candidates were: %s",
				file, candidatesStr)
		}

		return foundReq.SourceUnit(), nil
	}

	// Case 3: in std lib
	pythonDirIdx := -1
	for i, cmp := range fileCmps {
		if strings.HasPrefix(cmp, "python") {
			pythonDirIdx = i
			break
		}
	}
	if pythonDirIdx != -1 {
		return stdLibPkg.SourceUnit(), nil
	}

	return nil, fmt.Errorf("Cannot infer source unit for file %s", file)
}

func isSubPath(parent, child string) bool {
	relpath, err := filepath.Rel(parent, child)
	return err == nil && !strings.HasPrefix(relpath, "..")
}

func moduleToFilepath(moduleName string, isPackage bool) string {
	moduleName = strings.Replace(moduleName, ".", "/", -1)
	if !isPackage {
		moduleName += ".py"
	}
	return moduleName
}

type RawOutput struct {
	Defs []*RawDef
	Refs []*RawRef
}

type RawDef struct {
	Path      string
	Kind      string
	Name      string
	File      string // relative path (to source unit directory)
	DefStart  uint32
	DefEnd    uint32
	Exported  bool
	Docstring string
	Data      interface{}
}

type RawRef struct {
	DefPath   string
	Def       bool
	DefFile   string // absolute path
	File      string // relative path (to source unit directory)
	Start     uint32
	End       uint32
	ToBuiltin bool
}

func installPipRequirements(pipBin string, requirementFiles []string) {
	for _, requirementFile := range requirementFiles {
		err := runCmdStderr(exec.Command(pipBin, "install", "-r", requirementFile))
		if err != nil {
			log.Printf("Error installing dependencies in %s. Trying piecemeal install", requirementFile)
			if b, err := ioutil.ReadFile(requirementFile); err == nil {
				for _, req := range strings.Split(string(b), "\n") {
					runCmdLogError(exec.Command(pipBin, "install", req))
				}
			} else {
				log.Printf("Could not read %s: %s", requirementFile, err)
			}
		}
	}
}
<commit_msg>Speedup graphing.<commit_after>package python

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"

	"strings"

	"sourcegraph.com/sourcegraph/srclib/graph"
	"sourcegraph.com/sourcegraph/srclib/toolchain"
	"sourcegraph.com/sourcegraph/srclib/unit"
)

type GraphContext struct {
	Unit *unit.SourceUnit
	Reqs []*requirement
}

func NewGraphContext(unit *unit.SourceUnit) *GraphContext {
	var g GraphContext
	g.Unit = unit
	for _, dep := range unit.Dependencies {
		if req, err := asRequirement(dep); err == nil {
			g.Reqs = append(g.Reqs, req)
		}
	}
	return &g
}

// Graphs the Python source unit. If run outside of a Docker container, this assumes that the source unit has already
// been installed (via pip or `python setup.py install`).
func (c *GraphContext) Graph() (*graph.Output, error) {
	programMode := os.Getenv("IN_DOCKER_CONTAINER") == ""
	tc, err := toolchain.Lookup("sourcegraph.com/sourcegraph/srclib-python")
	if err != nil {
		return nil, err
	}
	pipBin := "pip"
	pythonBin := "python"

	if programMode {
		tempPath, err := getTempPath()
		if err != nil {
			return nil, err
		}
		envName := fmt.Sprintf("%s-%s-env", getHash(c.Unit.Dir), url.QueryEscape(c.Unit.Name))
		envDir := filepath.Join(tempPath, envName)

		// Use binaries from our virutal env.
		pipBin = filepath.Join(envDir, "bin", "pip")
		pythonBin = filepath.Join(envDir, "bin", "python")

		if _, err := os.Stat(filepath.Join(envDir)); os.IsNotExist(err) {
			// We don't have virtual env for this SourceUnit, create one.
			tcVENVBinPath := filepath.Join(tc.Dir, ".env", "bin")
			cmd := exec.Command(filepath.Join(tcVENVBinPath, "virtualenv"), envDir)
			if err := runCmdStderr(cmd); err != nil {
				return nil, err
			}
			// Install our dependencies.
			// Todo(MaikuMori): Use symlinks from toolchains virtualenv to project virtual env.
			// NOTE: If SourceUnit requirements overwrite our requirements, things will fail.
			// 			 We could install them last, but then we would have to do this before each
			//			 graphing which noticably increases graphing time (since our deps are always
			//       downloaded by pip due to dependency on git commit not actual package version).
			requirementFile := filepath.Join(tc.Dir, "requirements.txt")
			if err := runCmdStderr(exec.Command(pipBin, "install", "-r", requirementFile)); err != nil {
				return nil, err
			}
			if err := runCmdStderr(exec.Command(pipBin, "install", "-I", tc.Dir)); err != nil {
				return nil, err
			}
		}
	}

	// NOTE: this may cause an error when graphing any source unit that depends
	// on jedi (or any other dependency of the graph code)
	requirementFiles, err := filepath.Glob(filepath.Join(c.Unit.Dir, "*requirements*.txt"))
	if err != nil {
		return nil, err
	}
	if _, err := os.Stat(filepath.Join(c.Unit.Dir, "setup.py")); !os.IsNotExist(err) {
		runCmdLogError(exec.Command(pipBin, "install", "-I", c.Unit.Dir))
	}
	installPipRequirements(pipBin, requirementFiles)

	cmd := exec.Command(pythonBin, "-m", "grapher.graph", "--verbose", "--dir", c.Unit.Dir, "--files")
	cmd.Args = append(cmd.Args, c.Unit.Files...)
	cmd.Stderr = os.Stderr
	log.Printf("Running %v", cmd.Args)
	b, err := cmd.Output()
	if err != nil {
		return nil, err
	}

	var raw RawOutput
	if err := json.Unmarshal(b, &raw); err != nil {
		return nil, err
	}

	out := c.transform(&raw, c.Unit)
	return out, nil
}

func (c *GraphContext) transform(raw *RawOutput, unit *unit.SourceUnit) *graph.Output {
	var out graph.Output

	for _, def := range raw.Defs {
		out.Defs = append(out.Defs, c.transformDef(def))
		if doc := c.transformDefDoc(def); doc != nil {
			out.Docs = append(out.Docs, doc)
		}
	}
	for _, ref := range raw.Refs {
		if outRef, err := c.transformRef(ref); err == nil {
			out.Refs = append(out.Refs, outRef)
		} else {
			log.Printf("Could not transform ref %v: %s", ref, err)
		}
	}

	return &out
}

var jediKindToDefKind = map[string]string{
	"statement":        "var",
	"statementelement": "var",
	"param":            "var",
	"module":           "module",
	"submodule":        "module",
	"class":            "type",
	"function":         "func",
	"lambda":           "func",
	"import":           "var",
}

func (c *GraphContext) transformDef(rawDef *RawDef) *graph.Def {
	return &graph.Def{
		DefKey: graph.DefKey{
			Repo:     c.Unit.Repo,
			Unit:     c.Unit.Name,
			UnitType: c.Unit.Type,
			Path:     string(rawDef.Path),
		},
		TreePath: string(rawDef.Path), // TODO: make this consistent w/ old way
		Kind:     jediKindToDefKind[rawDef.Kind],
		Name:     rawDef.Name,
		File:     rawDef.File,
		DefStart: rawDef.DefStart,
		DefEnd:   rawDef.DefEnd,
		Exported: rawDef.Exported,
		Data:     nil, // TODO
	}
}

func (c *GraphContext) transformRef(rawRef *RawRef) (*graph.Ref, error) {
	defUnit, err := c.inferSourceUnit(rawRef, c.Reqs)
	if err != nil {
		return nil, err
	}

	defPath := string(rawRef.DefPath)
	if defPath == "" {
		defPath = "."
	}

	return &graph.Ref{
		DefRepo:     defUnit.Repo,
		DefUnitType: defUnit.Type,
		DefUnit:     defUnit.Name,
		DefPath:     defPath,

		Repo:     c.Unit.Repo,
		Unit:     c.Unit.Name,
		UnitType: c.Unit.Type,

		File:  rawRef.File,
		Start: rawRef.Start,
		End:   rawRef.End,
		Def:   rawRef.Def,
	}, nil
}

func (c *GraphContext) transformDefDoc(rawDef *RawDef) *graph.Doc {
	return nil
}

func (c *GraphContext) inferSourceUnit(rawRef *RawRef, reqs []*requirement) (*unit.SourceUnit, error) {
	if rawRef.ToBuiltin {
		return stdLibPkg.SourceUnit(), nil
	}
	return c.inferSourceUnitFromFile(rawRef.DefFile, reqs)
}

// Note: file is expected to be an absolute path
func (c *GraphContext) inferSourceUnitFromFile(file string, reqs []*requirement) (*unit.SourceUnit, error) {
	// Case: in current source unit (u)
	pwd, _ := os.Getwd()
	if isSubPath(pwd, file) {
		return c.Unit, nil
	}

	// Case: in dependent source unit(depUnits)
	fileCmps := strings.Split(file, string(filepath.Separator))
	pkgsDirIdx := -1
	for i, cmp := range fileCmps {
		if cmp == "site-packages" || cmp == "dist-packages" {
			pkgsDirIdx = i
			break
		}
	}
	if pkgsDirIdx != -1 {
		fileSubCmps := fileCmps[pkgsDirIdx+1:]
		fileSubPath := filepath.Join(fileSubCmps...)

		var foundReq *requirement
	FindReq:
		for _, req := range reqs {
			for _, pkg := range req.Packages {
				if isSubPath(moduleToFilepath(pkg, true), fileSubPath) {
					foundReq = req
					break FindReq
				}
			}
			for _, mod := range req.Modules {
				if moduleToFilepath(mod, false) == fileSubPath {
					foundReq = req
					break FindReq
				}
			}
		}

		if foundReq == nil {
			var candidatesStr string
			if len(reqs) <= 7 {
				candidatesStr = fmt.Sprintf("%v", reqs)
			} else {
				candidatesStr = fmt.Sprintf("%v...", reqs[:7])
			}
			// FIXME: This doesn't work, note the pointer in `[]*requirement`. As error you get
			//        string representation of array of pointers.
			return nil, fmt.Errorf("Could not find requirement that contains file %s. Candidates were: %s",
				file, candidatesStr)
		}

		return foundReq.SourceUnit(), nil
	}

	// Case 3: in std lib
	pythonDirIdx := -1
	for i, cmp := range fileCmps {
		if strings.HasPrefix(cmp, "python") {
			pythonDirIdx = i
			break
		}
	}
	if pythonDirIdx != -1 {
		return stdLibPkg.SourceUnit(), nil
	}

	return nil, fmt.Errorf("Cannot infer source unit for file %s", file)
}

func isSubPath(parent, child string) bool {
	relpath, err := filepath.Rel(parent, child)
	return err == nil && !strings.HasPrefix(relpath, "..")
}

func moduleToFilepath(moduleName string, isPackage bool) string {
	moduleName = strings.Replace(moduleName, ".", "/", -1)
	if !isPackage {
		moduleName += ".py"
	}
	return moduleName
}

type RawOutput struct {
	Defs []*RawDef
	Refs []*RawRef
}

type RawDef struct {
	Path      string
	Kind      string
	Name      string
	File      string // relative path (to source unit directory)
	DefStart  uint32
	DefEnd    uint32
	Exported  bool
	Docstring string
	Data      interface{}
}

type RawRef struct {
	DefPath   string
	Def       bool
	DefFile   string // absolute path
	File      string // relative path (to source unit directory)
	Start     uint32
	End       uint32
	ToBuiltin bool
}

func installPipRequirements(pipBin string, requirementFiles []string) {
	for _, requirementFile := range requirementFiles {
		err := runCmdStderr(exec.Command(pipBin, "install", "-r", requirementFile))
		if err != nil {
			log.Printf("Error installing dependencies in %s. Trying piecemeal install", requirementFile)
			if b, err := ioutil.ReadFile(requirementFile); err == nil {
				for _, req := range strings.Split(string(b), "\n") {
					runCmdLogError(exec.Command(pipBin, "install", req))
				}
			} else {
				log.Printf("Could not read %s: %s", requirementFile, err)
			}
		}
	}
}
<|endoftext|>