sc2ai 0.0.4 → 0.0.5

data/lib/sc2ai/player/units.rb

@@ -30,6 +30,53 @@ module Sc2
       #   @return [Sc2::UnitGroup] a group of neutral units
       attr_accessor :effects # not a unit
 
+      # Returns the upgrade ids you have acquired such as weapon upgrade and armor upgrade ids.
+      # Shorthand for observation.raw_data.player.upgrade_ids
+      # @!attribute [r] upgrades_completed
+      #   @return [Array<Integer>] a group of neutral units
+      def upgrades_completed = observation&.raw_data&.player&.upgrade_ids.to_a || [] # not a unit
+
+      # Returns the upgrade ids which are researching or queued
+      # Not set for enemy.
+      # @return [Array<Integer>]
+      def upgrades_in_progress
+        # We need to scan every structure which performs upgrades for any order with an upgrade ability
+
+        result = []
+        # Loop every upgrade structure
+        structures
+          .select_type(Api::TechTree.upgrade_structure_unit_type_ids)
+          .each do |structure|
+          next unless structure.is_active? # Skip idle
+
+          # Check if any order at a structure contains an upgrade ability
+          structure.orders.each do |order|
+            Api::TechTree.upgrade_ability_data(structure.unit_type).each do |upgrade_id, update_info|
+              if update_info[:ability] == order.ability_id
+                # Save the upgrade_id
+                result << upgrade_id
+              end
+            end
+          end
+        end
+
+        # If the API told use it's complete, but an order still lingers, trust the API
+        result - upgrades_completed
+      end
+
+      # Returns the upgrade ids which are researching or queued
+      # @return [Boolean]
+      def upgrade_in_progress?(upgrade_id)
+        structure_unit_type_id = Api::TechTree.upgrade_researched_from(upgrade_id: upgrade_id)
+        research_ability_id = Api::TechTree.upgrade_research_ability_id(upgrade_id: upgrade_id)
+        structures.select_type(structure_unit_type_id).each do |structure|
+          structure.orders.each do |order|
+            return true if order.ability_id == research_ability_id
+          end
+        end
+        false
+      end
+
       # An array of Protoss power sources, which have a point, radius and unit tag
       # @!attribute power_sources
       #   @return [Array<Api::PowerSource>] an array of power sources
@@ -107,6 +154,13 @@ module Sc2
         data.abilities[ability_id]
       end
 
+      # Returns static [Api::UpgradeData] for an upgrade id
+      # @param upgrade_id [Integer] Api::UpgradeId::*
+      # @return [Api::UpgradeData]
+      def upgrade_data(upgrade_id)
+        data.upgrades[upgrade_id]
+      end
+
       # Checks unit data for an attribute value
       # @param unit [Integer,Api::Unit] Api::UnitTypeId or Api::Unit
       # @param attribute [Symbol] Api::Attribute, i.e. Api::Attribute::Mechanical or :Mechanical
@@ -141,17 +195,13 @@ module Sc2
       def subtract_cost(unit_type_id)
         unit_type_data = unit_data(unit_type_id)
 
-        # food_required is a float. ensure half units are counted as full
-        # TODO: Extend UnitTypeData message. def food_required = unit_id ==  Api::UnitTypeId::ZERGLING ? 1 : send("method_missing", :food_required)
-        supply_cost = unit_type_data.food_required
-        supply_cost = 1 if unit_type_id == Api::UnitTypeId::ZERGLING
-
         @spent_minerals += unit_type_data.mineral_cost
         @spent_vespene += unit_type_data.vespene_cost
-        @spent_supply += supply_cost
+        @spent_supply += unit_type_data.food_required
       end
 
       # Checks whether you have the resources to construct quantity of unit type
+      # @return [Boolean]
       def can_afford?(unit_type_id:, quantity: 1)
         unit_type_data = unit_data(unit_type_id)
         return false if unit_type_data.nil?
@@ -178,6 +228,25 @@ module Sc2
         true
       end
 
+      # Checks whether you have the resources to
+      # @return [Boolean]
+      def can_afford_upgrade?(upgrade_id)
+        unit_type_data = upgrade_data(upgrade_id)
+        return false if unit_type_data.nil?
+
+        mineral_cost = unit_type_data.mineral_cost
+        if common.minerals - spent_minerals < mineral_cost
+          return false # not enough minerals
+        end
+
+        vespene_cost = unit_type_data.vespene_cost
+        if common.vespene - spent_vespene < vespene_cost
+          return false # you require more vespene gas
+        end
+
+        true
+      end
+
       private
 
       # @private

data/lib/sc2ai/protocol/_meta_documentation.rb

@@ -21,6 +21,8 @@
 #     class Point < Google::Protobuf::AbstractMessage; end;
 #     # Protobuf virtual class.
 #     class Unit < Google::Protobuf::AbstractMessage; end;
+#     # Protobuf virtual class.
+#     class UnitTypeData < Google::Protobuf::AbstractMessage; end;
 #   end
 
 # Protobuf enums ---

data/lib/sc2ai/protocol/extensions/point_2_d.rb

@@ -10,6 +10,12 @@ module Api
       self.class == other.class && hash == other.hash
     end
 
+    # Create a new 3d Point, by adding a y axis.
+    # @return [Api::Point]
+    def to_3d(z:)
+      Api::Point[x, y, z]
+    end
+
     # Adds additional functionality to message class Api::Point2D
     module ClassMethods
       # Shorthand for creating an instance for [x, y]

data/lib/sc2ai/protocol/extensions/unit.rb

@@ -37,9 +37,8 @@ module Api
     # Checks unit data for an attribute value
     # @return [Boolean] whether unit has attribute
     # @example
-    #   has_attribute?(Api::UnitTypeId::SCV, Api::Attribute::Mechanical)
-    #   has_attribute?(units.workers.first, :Mechanical)
-    #   has_attribute?(Api::UnitTypeId::SCV, :Mechanical)
+    #   unit.has_attribute?(Api::Attribute::Mechanical)
+    #   unit.has_attribute?(:Mechanical)
     def has_attribute?(attribute)
       attributes.include? attribute
     end
@@ -233,10 +232,18 @@ module Api
 
     # Issues repair command on target
     # @param target [Api::Unit, Integer] is a unit or unit tag
+    # @param queue_command [Boolean] shift+command
     def repair(target:, queue_command: false)
       action(ability_id: Api::AbilityId::EFFECT_REPAIR, target:, queue_command:)
     end
 
+    # Research a specific upgrade
+    # @param upgrade_id [Integer] Api::UnitTypeId the unit type which will do the creation
+    # @param queue_command [Boolean] shift+command
+    def research(upgrade_id:, queue_command: false)
+      @bot.research(units: self, upgrade_id:, queue_command:)
+    end
+
     # @!endgroup Actions
     #
     # Debug ----

data/lib/sc2ai/protocol/extensions/unit_type_data.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module Api
+  # Adds additional functionality to message object Api::UnitTypeData
+  module UnitTypeDataExtension
+    # @!attribute mineral_cost_sum
+    #   Sum of all morphs mineral cost
+    #   i.e. 550M Orbital command = 400M CC + 150M Upgrade
+    #   i.e. 350M Hatchery = 50M Drone + 300M Build
+    #   @return [Integer] sum of mineral costs
+    attr_accessor :mineral_cost_sum
+
+    # @!attribute vespene_cost_sum
+    #   Sum of all morphs vespene gas cost
+    #   i.e. 250G Broodlord = 100G Corruptor + 150G Morph
+    #   @return [Integer] sum of vespene gas costs
+    attr_accessor :vespene_cost_sum
+  end
+end
+Api::UnitTypeData.include Api::UnitTypeDataExtension

data/lib/sc2ai/unit_group/action_ext.rb

@@ -49,6 +49,15 @@ module Sc2
       bot&.warp(units: self, unit_type_id:, target:, queue_command:)
     end
 
+    # Research a specific upgrade at one of these structures
+    # @param upgrade_id [Integer] Api::UpgradeId to research
+    # @param queue_command [Boolean] shift+command
+    def research(upgrade_id:, queue_command: false)
+      return if size.zero?
+
+      bot&.research(units: self, upgrade_id:, queue_command:)
+    end
+
     # Shorthand for performing action SMART (right-click)
     # @param target [Api::Unit, Integer, Api::Point2D] is a unit, unit tag or a Api::Point2D
     # @param queue_command [Boolean] shift+command

data/lib/sc2ai/unit_group/filter_ext.rb

@@ -229,6 +229,12 @@ module Sc2
     alias_method :extractors, :gas
     alias_method :assimilators, :gas
 
+    # Selects only units which have finished constructing, i.o.w. build_progress == 1.0
+    # @return [UnitGroup] gas structures
+    def completed
+      select(&:is_completed?)
+    end
+
     # NEUTRAL ------------------------------------------
 
     # Selects mineral fields

data/lib/sc2ai/unit_group.rb

@@ -227,12 +227,20 @@ module Sc2
       UnitGroup.new(@units.except(...))
     end
 
-    # Returns a hash containing the entries for given tag(s).
+    # Returns a new unit group containing the entries for given tag(s).
     # @return [Sc2::UnitGroup] new unit group
     def slice(...)
       UnitGroup.new(@units.slice(...))
     end
 
+    # Returns a new unit group containing each element found both in self and in all of the given other_arrays; duplicates are omitted; items are compared using eql? (items must also implement hash correctly):
+    # @param other_unit_group [UnitGroup]
+    # @return [UnitGroup] new unit group
+    def intersection(other_unit_group)
+      slice(*other_unit_group.tags)
+    end
+    alias_method :&, :intersection
+
     # Selects a single random Unit without a parameter or an array of Units with a param, i.e. self.random(2)
     # @return [Api::Unit]
     def sample(...)

data/lib/sc2ai/version.rb

@@ -2,5 +2,5 @@
 
 module Sc2
   # gem version
-  VERSION = "0.0.4"
+  VERSION = "0.0.5"
 end

data/lib/templates/new/run_example_match.rb.tt

@@ -10,5 +10,5 @@ Sc2::Match.new(
     $bot,
     Sc2::Player::Computer.new(race: Api::Race::Random, difficulty: Api::Difficulty::VeryEasy)
   ],
-  map: "Goldenaura512AIE" # Or any of the downloaded map names
+  map: "Goldenaura512V2AIE" # Or any of the downloaded map names
 ).run