Add getters for services that have defined build, extends, depends_on, capabilities, gpu and tpu

types/project.go

@@ -29,6 +29,7 @@ import (
 	"github.com/distribution/reference"
 	"github.com/mitchellh/copystructure"
 	godigest "github.com/opencontainers/go-digest"
+	"golang.org/x/exp/maps"
 	"golang.org/x/sync/errgroup"
 	"gopkg.in/yaml.v3"
 )
@@ -119,6 +120,58 @@ func (p *Project) ConfigNames() []string {
 	return names
 }
 
+func (p *Project) ServicesBuild() []string {
+	servicesBuild := p.Services.Filter(func(s ServiceConfig) bool {
+		return s.Build != nil && s.Build.Context != ""
+	})
+	return maps.Keys(servicesBuild)
+}
+
+func (p *Project) ServicesWithExtends() []string {
+	servicesExtends := p.Services.Filter(func(s ServiceConfig) bool {
+		return s.Extends != nil && *s.Extends != (ExtendsConfig{})
+	})
+	return maps.Keys(servicesExtends)
+}
+
+func (p *Project) ServicesWithDependsOn() []string {
+	servicesDependsOn := p.Services.Filter(func(s ServiceConfig) bool {
+		return len(s.DependsOn) > 0
+	})
+	return maps.Keys(servicesDependsOn)
+}
+
+func (p *Project) ServicesWithCapabilities() ([]string, []string, []string) {
+	capabilities := []string{}
+	gpu := []string{}
+	tpu := []string{}
+	for _, service := range p.Services {
+		deploy := service.Deploy
+		if deploy == nil {
+			continue
+		}
+		reservation := deploy.Resources.Reservations
+		if reservation == nil {
+			continue
+		}
+		devices := reservation.Devices
+		for _, d := range devices {
+			if len(d.Capabilities) > 0 {
+				capabilities = append(capabilities, service.Name)
+			}
+			for _, c := range d.Capabilities {
+				if c == "gpu" {
+					gpu = append(gpu, service.Name)
+				} else if c == "tpu" {
+					tpu = append(tpu, service.Name)
+				}
+			}
+		}
+	}
+
+	return utils.RemoveDuplicates(capabilities), utils.RemoveDuplicates(gpu), utils.RemoveDuplicates(tpu)
+}
+
 // GetServices retrieve services by names, or return all services if no name specified
 func (p *Project) GetServices(names ...string) (Services, error) {
 	if len(names) == 0 {

types/project_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/compose-spec/compose-go/v2/utils"
 	"github.com/distribution/reference"
 	"github.com/opencontainers/go-digest"
+	"golang.org/x/exp/slices"
 	"gotest.tools/v3/assert"
 )
 
@@ -240,3 +241,133 @@ func TestWithServices(t *testing.T) {
 	assert.NilError(t, err)
 	assert.DeepEqual(t, seen, []string{"service_1", "service_2", "service_4"})
 }
+
+func TestBuildServices(t *testing.T) {
+	p := makeProject()
+	assert.DeepEqual(t, []string{}, p.ServicesBuild())
+
+	service, err := p.GetService("service_1")
+	assert.NilError(t, err)
+	service.Build = &BuildConfig{}
+	p.Services["service_1"] = service
+	assert.DeepEqual(t, []string{}, p.ServicesBuild())
+
+	service.Build = &BuildConfig{
+		Context: ".",
+	}
+	p.Services["service_1"] = service
+	service, err = p.GetService("service_2")
+	assert.NilError(t, err)
+	service.Build = &BuildConfig{
+		Context: ".",
+	}
+	p.Services["service_2"] = service
+	services := p.ServicesBuild()
+	slices.Sort(services)
+	assert.DeepEqual(t, []string{"service_1", "service_2"}, services)
+}
+
+func TestServicesWithExtends(t *testing.T) {
+	p := makeProject()
+	assert.DeepEqual(t, []string{}, p.ServicesWithExtends())
+
+	service, err := p.GetService("service_1")
+	assert.NilError(t, err)
+	service.Extends = &ExtendsConfig{}
+	p.Services["service_1"] = service
+	assert.DeepEqual(t, []string{}, p.ServicesWithExtends())
+
+	service.Extends = &ExtendsConfig{
+		File:    ".",
+		Service: "service_2",
+	}
+	p.Services["service_1"] = service
+	assert.DeepEqual(t, []string{"service_1"}, p.ServicesWithExtends())
+}
+
+func TestServicesWithDependsOn(t *testing.T) {
+	p := makeProject()
+	services := p.ServicesWithDependsOn()
+	slices.Sort(services)
+	assert.Equal(t, 3, len(services))
+	assert.DeepEqual(t, []string{"service_2", "service_3", "service_4"}, services)
+}
+
+func TestServicesWithCapabilities(t *testing.T) {
+	p := makeProject()
+
+	service, err := p.GetService("service_1")
+	assert.NilError(t, err)
+	service.Deploy = &DeployConfig{}
+	p.Services["service_1"] = service
+	capabilities, gpu, tpu := p.ServicesWithCapabilities()
+	assert.DeepEqual(t, []string{}, capabilities)
+	assert.DeepEqual(t, []string{}, gpu)
+	assert.DeepEqual(t, []string{}, tpu)
+
+	service.Deploy = &DeployConfig{
+		Resources: Resources{
+			Reservations: &Resource{},
+		},
+	}
+	p.Services["service_1"] = service
+	capabilities, gpu, tpu = p.ServicesWithCapabilities()
+	assert.DeepEqual(t, []string{}, capabilities)
+	assert.DeepEqual(t, []string{}, gpu)
+	assert.DeepEqual(t, []string{}, tpu)
+
+	service.Deploy = &DeployConfig{
+		Resources: Resources{
+			Reservations: &Resource{
+				Devices: []DeviceRequest(nil),
+			},
+		},
+	}
+	p.Services["service_1"] = service
+	capabilities, gpu, tpu = p.ServicesWithCapabilities()
+	assert.DeepEqual(t, []string{}, capabilities)
+	assert.DeepEqual(t, []string{}, gpu)
+	assert.DeepEqual(t, []string{}, tpu)
+
+	service.Deploy = &DeployConfig{
+		Resources: Resources{
+			Reservations: &Resource{
+				Devices: []DeviceRequest{
+					{
+						Capabilities: []string{"gpu", "tpu"},
+					},
+				},
+			},
+		},
+	}
+	p.Services["service_1"] = service
+	capabilities, gpu, tpu = p.ServicesWithCapabilities()
+	assert.DeepEqual(t, []string{"service_1"}, capabilities)
+	assert.DeepEqual(t, []string{"service_1"}, gpu)
+	assert.DeepEqual(t, []string{"service_1"}, tpu)
+
+	service, err = p.GetService("service_2")
+	assert.NilError(t, err)
+	service.Deploy = &DeployConfig{
+		Resources: Resources{
+			Reservations: &Resource{
+				Devices: []DeviceRequest{
+					{
+						Capabilities: []string{"tpu"},
+					},
+					{
+						Capabilities: []string{"tpu"},
+					},
+				},
+			},
+		},
+	}
+	p.Services["service_2"] = service
+	capabilities, gpu, tpu = p.ServicesWithCapabilities()
+	slices.Sort(capabilities)
+	slices.Sort(gpu)
+	slices.Sort(tpu)
+	assert.DeepEqual(t, []string{"service_1", "service_2"}, capabilities)
+	assert.DeepEqual(t, []string{"service_1"}, gpu)
+	assert.DeepEqual(t, []string{"service_1", "service_2"}, tpu)
+}

types/services.go

@@ -33,3 +33,13 @@ func (s Services) GetProfiles() []string {
 	}
 	return profiles
 }
+
+func (s Services) Filter(predicate func(ServiceConfig) bool) Services {
+	services := Services{}
+	for name, service := range s {
+		if predicate(service) {
+			services[name] = service
+		}
+	}
+	return services
+}

utils/collectionutils.go

@@ -51,3 +51,18 @@ func ArrayContains[T comparable](source []T, toCheck []T) bool {
 	}
 	return true
 }
+
+func RemoveDuplicates[T comparable](slice []T) []T {
+	// Create a map to store unique elements
+	seen := make(map[T]bool)
+	result := []T{}
+
+	// Loop through the slice, adding elements to the map if they haven't been seen before
+	for _, val := range slice {
+		if _, ok := seen[val]; !ok {
+			seen[val] = true
+			result = append(result, val)
+		}
+	}
+	return result
+}